OREO: Offline REasoning Optimization

Implementation of OREO (Offline REasoning Optimization) - an offline reinforcement learning method for enhancing large language models' multi-step reasoning capabilities.

📚 What is OREO?

OREO is an offline RL algorithm that improves LLM reasoning by:

• Jointly learning a policy model (LLM) and value function
• Optimizing the soft Bellman equation for better credit assignment
• Reducing reliance on paired preference data (vs DPO)
• Handling sparse rewards better through value-based bootstrapping

Paper: Offline Reinforcement Learning for LLM Multi-Step Reasoning Original Code: https://github.com/jwhj/oreo

---

🚀 Quick Start

1. Installation

# Install OREO-specific dependencies
pip install -r requirements_oreo.txt

# Alternatively, install individually:
pip install transformers>=4.36.0 datasets>=2.14.0 accelerate>=0.25.0

2. Download Data

# Download GSM8K dataset
python dataset/download_data.py --dataset gsm8k

# Or use the script (Windows)
scripts\01_download_data.bat

3. Collect Trajectories (Optional)

For full OREO training, collect reasoning trajectories:

# Collect trajectories (takes 15-30 mins for 100 problems)
python -m dataset.trajectory_collector \
  --base_model "Qwen/Qwen2.5-Math-1.5B" \
  --data_file dataset/gsm8k_train.jsonl \
  --output_file dataset/trajectories_gsm8k.jsonl \
  --num_trajectories_per_problem 4 \
  --max_problems 100

# Or use the script (Windows)
scripts\02_collect_trajectories.bat

4. Train OREO

Phase 1: Supervised Fine-Tuning

python main.py \
  --mode oreo_sft \
  --base_model "Qwen/Qwen2.5-Math-1.5B" \
  --reasoning_dataset gsm8k \
  --epochs 3 \
  --batch_size 4 \
  --lr 1e-5

# Or use the script (Windows)
scripts\03_train_sft.bat

Phase 2: OREO Offline RL

python main.py \
  --mode oreo_rl \
  --base_model "Qwen/Qwen2.5-Math-1.5B" \
  --trajectory_file dataset/trajectories_gsm8k.jsonl \
  --epochs 10 \
  --batch_size 4 \
  --policy_lr 3e-6 \
  --value_lr 1e-5

# Or use the script (Windows)
scripts\04_train_oreo_rl.bat

5. Evaluate

python main.py \
  --mode test \
  --base_model "Qwen/Qwen2.5-Math-1.5B" \
  --ckpt_name oreo_policy_e10 \
  --reasoning_dataset gsm8k

# Or use the script (Windows)
scripts\05_evaluate.bat

---

📊 Complete Training Pipeline

# Step 0: Download data
python dataset/download_data.py --dataset gsm8k

# Step 1: Collect trajectories
python -m dataset.trajectory_collector \
  --base_model "Qwen/Qwen2.5-Math-1.5B" \
  --data_file dataset/gsm8k_train.jsonl \
  --output_file dataset/trajectories_gsm8k.jsonl \
  --max_problems 100

# Step 2: SFT training
python main.py --mode oreo_sft \
  --base_model "Qwen/Qwen2.5-Math-1.5B" \
  --reasoning_dataset gsm8k \
  --epochs 3 --batch_size 4

# Step 3: OREO RL training
python main.py --mode oreo_rl \
  --trajectory_file dataset/trajectories_gsm8k.jsonl \
  --epochs 10 --batch_size 4

# Step 4: Evaluation
python main.py --mode test \
  --ckpt_name oreo_policy_e10 \
  --reasoning_dataset gsm8k

---

🔧 Configuration

Key Hyperparameters

Memory Optimization for GTX 1060 (6GB)

# Use these settings for small GPU:
--batch_size 2 \
--use_gradient_checkpointing true \
--max_seq_length 256 \
--precision float16  # If supported

Multi-GPU Training

# Using torchrun (recommended)
torchrun --nproc_per_node=2 main.py \
  --distributed true \
  --mode oreo_sft \
  ...

# Effective batch size = batch_size * num_gpus

---

📁 Directory Structure

OREO/
├── models/OREO/              # OREO models
│   ├── policy_model.py       # Policy (LLM)
│   ├── value_model.py        # Value function
│   └── oreo_config.py        # Configuration
├── dataset/                  # Data handling
│   ├── download_data.py      # Download datasets
│   ├── trajectory_collector.py  # Collect trajectories
│   └── oreo_data_provider.py    # Data loading
├── learning/                 # Training logic
│   ├── oreo_task.py          # OREO task class
│   └── oreo_engine.py        # Training loops
├── utils/                    # Utilities
│   ├── oreo_losses.py        # Loss functions
│   ├── oreo_metrics.py       # Evaluation metrics
│   ├── generation_utils.py   # Text generation
│   └── answer_extraction.py  # Answer parsing
├── scripts/                  # Training scripts
│   ├── 01_download_data.bat
│   ├── 02_collect_trajectories.bat
│   ├── 03_train_sft.bat
│   ├── 04_train_oreo_rl.bat
│   └── 05_evaluate.bat
└── checkpoints/              # Saved models (created automatically)

---

🎯 Datasets

Supported Datasets

1. GSM8K (Grade School Math, 8.5K problems)

   • Simple arithmetic and word problems
   • Good for initial testing
   • Download: python dataset/download_data.py --dataset gsm8k

2. MATH (Competition Math)

   • More challenging problems
   • Download: python dataset/download_data.py --dataset math

3. Test (Small synthetic dataset)

   • 10 simple math problems
   • For quick testing
   • Download: python dataset/download_data.py --dataset test

Data Format

Trajectories are stored as JSONL with format:

{
  "problem": "What is 25% of 80?",
  "solution": "25% = 0.25. 0.25 * 80 = 20",
  "predicted_answer": "20",
  "ground_truth": "20",
  "is_correct": true,
  "reward": 1.0
}

---

🧠 Model Architecture

Policy Model

• Base: HuggingFace transformer (e.g., Qwen2.5-Math-1.5B)
• Purpose: Generate reasoning steps autoregressively
• Input: Problem + previous reasoning steps
• Output: Next reasoning step (token logits)

Value Model

• Base: Shared or separate transformer
• Purpose: Estimate value V(s) of reasoning states
• Architecture: Transformer + MLP value head
• Output: Scalar value estimate

Shared Backbone (Recommended)

• Policy and value share transformer weights
• More memory efficient
• Faster training
• Enable with --share_backbone true

---

📈 Training Tips

For Small GPUs (6GB)

--batch_size 2 \
--use_gradient_checkpointing true \
--max_seq_length 256 \
--max_problems 50  # For trajectory collection

For Faster Training

--batch_size 8 \
--epochs 5 \
--max_problems 50  # Fewer trajectories

For Best Accuracy

--epochs 20 \
--policy_lr 1e-6 \  # Lower learning rate
--temperature 0.05 \  # Lower temperature
--max_problems 500  # More trajectories

---

🔍 Troubleshooting

Out of Memory (OOM)

Solution 1: Reduce batch size

--batch_size 1  # or 2

Solution 2: Enable gradient checkpointing

--use_gradient_checkpointing true

Solution 3: Reduce sequence length

--max_seq_length 256  # or 128

Solution 4: Use smaller model

--base_model "Qwen/Qwen2-0.5B"  # 500M params instead of 1.5B

Slow Training

Solution 1: Use multi-GPU

torchrun --nproc_per_node=2 main.py --distributed true ...

Solution 2: Reduce data size

--max_problems 50  # Fewer trajectory examples

Solution 3: Fewer epochs

--epochs 5  # Instead of 10-20

Poor Accuracy

Solution 1: Train longer

--epochs 20  # More epochs

Solution 2: More trajectories

--max_problems 500  # More training data

Solution 3: Adjust hyperparameters

--policy_lr 1e-6 \  # Lower LR
--temperature 0.05  # Lower temperature

Generation Issues

Solution: Adjust generation parameters

--temperature 0.7 \  # Try 0.5-1.0
--top_p 0.9 \        # Try 0.8-0.95
--top_k 50           # Try 20-100

---

📊 Evaluation Metrics

• Accuracy: Final answer correctness (exact match)
• Average Reward: Mean reward over test set
• Trajectory Length: Average number of reasoning steps

Example output:

==================================================
EVALUATION RESULTS
==================================================
  accuracy: 0.7500
  avg_reward: 0.7500
  avg_trajectory_length: 3.2
  num_predictions: 100
==================================================

---

🆚 Comparison with Baselines

---

🔬 Advanced Usage

Custom Dataset

1. Create JSONL file with format:

{"problem": "...", "solution": "...", "answer": "..."}

2. Modify oreo_data_provider.py to load your dataset

3. Train with --reasoning_dataset custom

Different Base Model

--base_model "meta-llama/Llama-2-7b" \  # LLaMA
--base_model "deepseek-ai/deepseek-math-7b-instruct"  # DeepSeek

Hyperparameter Tuning

Key hyperparameters to tune:

• policy_lr: Try 1e-6 to 1e-5
• value_lr: Try 3e-6 to 3e-5
• temperature: Try 0.05 to 0.2
• gamma: Try 0.95 to 0.99

---

📚 References

Paper: Offline Reinforcement Learning for LLM Multi-Step Reasoning

• arXiv: https://arxiv.org/abs/2412.16145
• HuggingFace: https://huggingface.co/papers/2412.16145

Original Implementation:

• GitHub: https://github.com/jwhj/oreo

Pretrained Models:

• Policy: jwhj/Qwen2.5-Math-1.5B-OREO
• Value: jwhj/Qwen2.5-Math-1.5B-OREO-Value

---

🙋 FAQ

Q: Do I need to collect trajectories? A: For full OREO training (RL phase), yes. For SFT only, no.

Q: How long does training take? A: SFT: 30-60 min, RL: 1-2 hours (GTX 1060, batch_size=2, 100 problems)

Q: Can I use my own dataset? A: Yes! Create JSONL with problem/solution/answer fields.

Q: What GPU do I need? A: Minimum 6GB (GTX 1060). Recommended: 8GB+ (RTX 3060).

Q: Does it work on CPU? A: Yes, but very slow (not recommended).

Q: Can I skip SFT and train RL directly? A: Not recommended. SFT provides good initialization for RL.

---

✅ Implementation Checklist

All OREO components have been implemented:

• Policy model (HuggingFace LLM wrapper)
• Value model (transformer + value head)
• Soft Bellman loss function
• Value TD loss function
• SFT training loop
• OREO RL training loop
• Trajectory collection
• Data download scripts
• Evaluation metrics
• Answer extraction
• Generation utilities
• Multi-GPU (DDP) support
• Gradient checkpointing
• Training scripts
• Complete documentation

---

🎉 Success!

You're now ready to train OREO models for multi-step reasoning! Start with the Quick Start section above.

For questions or issues, refer to:

• This README
• OREO_IMPLEMENTATION_PLAN.md (detailed technical plan)
• Original OREO paper and code

Happy reasoning! 🧠✨