ExpRL: Exploratory RL for LLM Mid-Training

Official code for the paper "ExpRL: Exploratory RL for LLM Mid-Training" (paper).

Sparse-reward RL for LLM reasoning works only when the base model already has coverage over productive reasoning paths. ExpRL is an RL-based mid-training (priming) method that builds this coverage from large corpora of question–answer data. Rather than imitating reference solutions, ExpRL uses them as reward scaffolds: the reference is hidden from the policy and used only to let an LLM judge score on-policy reasoning traces for partial progress. Optimizing these dense, reference-guided rewards shifts probability mass toward productive reasoning paths, yielding a stronger initialization for subsequent sparse-reward RL than SFT, sparse-reward GRPO, or self-distillation.

The method has two stages, each implemented on a separate training stack that is vendored into this repo:

Base policy: Qwen3-4B-Instruct-2507 (also used as the judge). The mixed-domain experiments use an 8B policy with a 4B judge. Held-out benchmarks: AIME 2025, AIME 2026, HMMT (Nov 2025), IMO-AnswerBench (+ GPQA / Olympiad-Physics / LCB for mixed-domain).

ExpRL/
├── paper.pdf                 # the paper
├── stage1_verl/              # Stage 1: dense reference-guided RL priming (verl fork)
│   └── variational_inference/math/exp_rl/   # entry-point scripts
│       ├── train_stage1/     #   priming runs (ExpRL-Outcome/-Process + baselines)
│       ├── inference_stage1/ #   sample rollouts for pass@k on held-out benchmarks
│       ├── inference_stage2/ #   sample rollouts from Stage-2 checkpoints
│       └── eval/             #   pass@1 / pass@k scoring
└── stage2_pipeline_rl/       # Stage 2: downstream sparse-reward GRPO (PipelineRL fork)
    ├── conf/                 #   exp_stage2_*.yaml configs
    └── scripts/              #   run_exp_stage2_*.sh entry points

Each stage has its own README with detailed setup and run commands: stage1_verl/README.md, stage2_pipeline_rl/README.md.

The two stages use separate conda environments (different, partly conflicting pins for vLLM / torch / flash-attn).

Stage 1 (verl):

conda create -n verl python=3.10 -y && conda activate verl
cd stage1_verl && pip install -e .
# newer torchao breaks the transformers import on torch 2.6:
pip install torchao==0.9.0

Stage 2 (prl): needs a CUDA 12.4 toolkit with CUDA_HOME set (flash-attn / deepspeed build from source):

conda create -n prl python=3.11 -y && conda activate prl
cd stage2_pipeline_rl && ./install.sh

The Stage-2 pyproject.toml pins fastapi==0.115.6 (newer starlette breaks the prometheus instrumentation in vLLM's OpenAI server).

Set the usual credentials yourself: HF_TOKEN (model/dataset access), optionally WANDB_API_KEY (logging) and OPENAI_API_KEY / OPENAI_BASE_URL (if you point the judge at an external endpoint).

The entry-point scripts are parametrized with two environment variables so they run on any machine without editing:

• EXPRL_ROOT — absolute path to the stage checkout (Stage-1 scripts cd here). Defaults to auto-detecting the repo root from the script location; set it explicitly when launching via SLURM: export EXPRL_ROOT=/path/to/ExpRL/stage1_verl.
• HF_NAMESPACE — Hugging Face org/user for checkpoints and intermediate datasets (defaults to violetxi, the authors' namespace). Set it to your own account when you push your Stage-1 outputs: export HF_NAMESPACE=<your-hf-username>.

Stage-2 configs additionally take a model_path override on the command line (the Stage-1 checkpoint to continue from), e.g. bash scripts/run_exp_stage2_grpo.sh model_path=/path/to/your/stage1_ckpt.

The violetxi/* names are the authors' artifacts. Either set HF_NAMESPACE to use your own re-trained checkpoints, or use the public violetxi/* checkpoints once released. SLURM scripts ship with placeholder --account=YOUR_ACCOUNT / --partition=YOUR_PARTITION.

1. Prepare data (Stage 1). The training scripts read data/instruct/*.parquet built from the math mid-training sources (InT + POPE; plus SciKnow + LCB for mixed-domain). See stage1_verl/README.md and variational_inference/math/data_processing/instruct.py.

2. Stage 1 — prime the base model (env verl, single node 8×H100):

   cd stage1_verl
   bash variational_inference/math/exp_rl/train_stage1/qwen4b/dense_outcome.sh   # ExpRL-Outcome
   bash variational_inference/math/exp_rl/train_stage1/qwen4b/dense_pr_delta.sh  # ExpRL-Process
   # baselines: grpo.sh, self_distill_lr1e-7.sh, sft.sh

3. Stage 2 — downstream sparse-reward RL from a Stage-1 checkpoint (env prl, single node 8×H100):

   cd stage2_pipeline_rl
   export HF_NAMESPACE=<your-hf-username>   # where your Stage-1 ckpts live
   bash scripts/run_exp_stage2_grpo.sh      # or dense_outcome / dense_process / sft / self_distill

4. Evaluate — sample rollouts (inference_stage1/, inference_stage2/) and score with eval/ (verl.trainer.main_eval) to get pass@1 / pass@k on the held-out benchmarks.

The dense-reward advantage estimators (GPPO, GPR_E, PR_BN, PR_DELTA) live in stage1_verl/verl/trainer/ppo/core_algos.py; the LLM-judge prompts in stage1_verl/verl/utils/prompt_templates/math_likelihood_{judge,process}.py. The 8B mixed-domain runs are under train_stage1/qwen8b/.

@article{xiang2026exprl,
  title  = {ExpRL: Exploratory RL for LLM Mid-Training},
  author = {Xiang, Violet and Setlur, Amrith and Blagden, Chase and Haber, Nick and Kumar, Aviral},
  year   = {2026}
}

This repo vendors forks of verl and PipelineRL; their upstream licenses and NOTICE/Notice.txt files are retained inside each stage directory. See LICENSE for this repository's license.