Dual-Pool Hybrid Reinforcement Learning for LLM Mathematical Reasoning
LUFFY-POOL extends the LUFFY framework with a dual-pool training architecture that simultaneously leverages on-policy and off-policy data for more efficient reinforcement learning of large language models. The system is designed for training mathematical reasoning capabilities using GRPO (Group Relative Policy Optimization) with importance-weighted off-policy corrections.
- Dual-Pool Training: Combines on-policy rollouts (fresh samples from the current policy) with off-policy data (pre-collected trajectories) in a unified training loop, using importance sampling to correct distribution mismatch.
- GRPO with Off-Policy Extension: Extends Group Relative Policy Optimization to support mixed-source advantage estimation, enabling more sample-efficient training.
- Distributed Training Pipeline: Built on VERL (Volcano Engine RL) with FSDP + Ray + vLLM, supporting 3D parallelism (tensor/pipeline/data) for large-scale training.
- Multi-Strategy Reward Functions: Supports rule-based symbolic verification,
math-verifylibrary grading, and LLM-as-Reward-Model (ORM) for nuanced evaluation of reasoning chains.
┌─────────────────────────────────────────────────┐
│ Mix Trainer │
│ │
│ ┌──────────────┐ ┌──────────────────┐ │
│ │ On-Policy │ │ Off-Policy │ │
│ │ Pool │ │ Pool │ │
│ │ (vLLM │ │ (Pre-collected │ │
│ │ Rollout) │ │ Trajectories) │ │
│ └──────┬───────┘ └────────┬──────────┘ │
│ │ │ │
│ └────────┬───────────────┘ │
│ ▼ │
│ ┌──────────────────────────────┐ │
│ │ GRPO Advantage Estimation │ │
│ │ (importance-weighted mix) │ │
│ └──────────────┬───────────────┘ │
│ ▼ │
│ ┌──────────────────────────────┐ │
│ │ FSDP Actor Update │ │
│ │ (gradient checkpointing) │ │
│ └──────────────────────────────┘ │
└─────────────────────────────────────────────────┘
LUFFY-POOL/
├── luffy/
│ ├── verl/ # VERL framework (Volcano Engine RL)
│ │ └── verl/mix_src/ # Core dual-pool implementation
│ │ ├── main_mix_ppo.py # Training entry point
│ │ ├── mix_trainer.py # Dual-pool trainer with on/off-policy learning
│ │ ├── mix_core_alg.py # GRPO advantage + on/off-policy loss functions
│ │ ├── mix_actor.py # Actor model training logic
│ │ ├── mix_vllm_rollout.py # vLLM-based generation and rollout
│ │ ├── mix_fsdp_worker.py # FSDP distributed training worker
│ │ └── rl_dataset_with_target.py # Dataset with target action sampling
│ ├── deepscaler/ # Reward module
│ │ ├── rewards/
│ │ │ ├── math_reward.py # Rule-based + LLM-as-RM reward functions
│ │ │ └── math_utils/ # Symbolic/numeric answer grading
│ │ └── system_prompts.py # ORM prompt templates
│ ├── requirements.txt
│ └── setup.py
├── data/ # Data preparation scripts & datasets
│ ├── prepare_sft.py # SFT data formatting
│ ├── prepare_train.py # RL training data
│ └── prepare_train_sft_rl.py # Mixed SFT+RL data
├── eval_scripts/ # Evaluation tools
│ ├── generate_vllm.py # vLLM batch generation
│ ├── collect_results.py # Result aggregation
│ └── oat_math_grader.py # Math grading utility
├── exp_scripts/ # Experiment configs (Hydra YAML)
└── figures/ # Visualizations
- On-Policy Rollout: The current policy generates multiple candidate solutions (n=8) per math problem using vLLM, forming the on-policy pool.
- Off-Policy Sampling: Pre-collected trajectories (from stronger models or previous checkpoints) are sampled with importance weights to form the off-policy pool.
- Mixed Advantage Estimation: GRPO computes group-relative advantages across both pools, with importance sampling ratios correcting for the off-policy distribution shift.
- Policy Update: The actor is updated via clipped PPO-style objectives on the combined loss, with KL penalty against a reference policy.
| Reward Type | Description |
|---|---|
| Rule-based verification | Symbolic and numeric answer matching with tolerance |
math-verify |
Library-based formal verification of mathematical expressions |
| LLM-as-RM (ORM) | GPT/Gemini evaluates reasoning chain quality beyond final answer correctness |
| Format reward | Checks for proper <think>...</think> chain-of-thought formatting |
- Base Model: Qwen2.5-Math-7B (16k context, think variant)
- Rollout: vLLM with sampling, 8 candidates per prompt, up to 8192 response tokens
- Optimization: FSDP with gradient checkpointing, AdamW optimizer
- Batch Size: 128 (train), 512 (validation)
- Epochs: 30
- Orchestration: Ray for multi-node coordination
# 1. Install dependencies
cd luffy && pip install -e . && pip install -r requirements.txt
# 2. Prepare training data
python data/prepare_train.py
# 3. Launch training (see exp_scripts/ for full configs)
python luffy/verl/verl/mix_src/main_mix_ppo.py \
--config-path ../exp_scripts \
--config-name config.yaml
# 4. Evaluate
python eval_scripts/generate_vllm.py \
--model_path <checkpoint_path> \
--data_path data/valid.parquet- Python 3.9+
- PyTorch 2.4.0
- Transformers 4.46.3
- vLLM 0.6.3
- Ray 2.12.0
- DeepSpeed 0.15.0
- Flash Attention 2.7.3
Built upon the LUFFY framework and VERL (Volcano Engine Reinforcement Learning).