Why Tree-Style Branching Matters for Thought Advantage Estimation in GRPO

Official implementation of our ICML 2026 paper Why Tree-Style Branching Matters for Thought Advantage Estimation in GRPO.

Hongcheng Wang*, Yinuo Huang*, Sukai Wang, Guanghui Ren†, Hao Dong☎   (*equal contribution  †project lead  ☎corresponding author)

Peking University  ·  PKU–Agibot Joint Lab  ·  UESTC  ·  Agibot

📄 Paper: https://arxiv.org/abs/2509.24494  |  🐙 Repo: https://github.com/whcpumpkin/GRPO-MA  |  ✉ Contact: Hongcheng Wang (whcpumpkin@gmail.com), Hao Dong (hao.dong@pku.edu.cn)

We cast a tree-style branching point as splitting a generation into a shared context and one or more continuations, and apply the multivariate delta method to the shared-context advantage estimator under GRPO-style group normalization. This reveals a sampling-dimension asymmetry: scaling thoughts K alone leaves a strictly positive variance floor, whereas scaling continuations M per shared context drives the leading-order estimation variance to zero at rate 1/M. Continuation-level branching is therefore principled and potentially necessary, not a heuristic. This bound is an algebraic invariance of the estimator, so it applies to any semantic breakpoint — not just the thought–answer split.

GRPO-MA is the lightweight instantiation we use to test the theory: place the breakpoint at </think> / </analysis> and sample M post-thought continuations per thought, with otherwise minimal changes to standard GRPO.

   prompt
     │
     ├── thought_1 ──► answer_{1,1}, answer_{1,2}, …, answer_{1,M}
     ├── thought_2 ──► answer_{2,1}, …
     │
     └── thought_K ──► answer_{K,1}, …, answer_{K,M}

This is a reference implementation of the GRPO-MA training framework. The paper evaluates eight task families; this v0.1.0 release ships the framework plus one canonical reference task (ShareRobot trajectory prediction). Other tasks are planned follow-ups.

The training framework is task-agnostic; new tasks plug in via task/TEMPLATE.py (see Adding a new task).

• Installation
• Quick start
• Project structure
• Training
• Evaluation
• Adding a new task
• Reproducibility notes
• Citation
• Acknowledgements
• License

git clone https://github.com/whcpumpkin/GRPO-MA.git
cd GRPO-MA
conda create -n grpo-ma python=3.10
conda activate grpo-ma
pip install -r requirements.txt
pip install -e .
# Optional: flash-attn for faster training
# pip install flash-attn==2.7.4 --no-build-isolation

The pinned versions in requirements.txt reflect the exact setup used for the paper experiments (PyTorch 2.5.1 + CUDA 12.4, transformers==4.51.3, trl==0.19.0, peft==0.15.2). A different toolchain may work but is not validated.

Sanity check:

python scripts/smoke_test.py

This verifies the task registry, reward routing, and metric dispatch without touching a GPU or downloading any model weights.

mkdir data && cd data
git clone https://huggingface.co/datasets/BAAI/ShareRobot
cd ..

To keep this repo small, the JSON manifests live in a companion HuggingFace dataset (see metadata/README.md for details and schema):

huggingface-cli download whcpumpkin/GRPO-MA-data \
    --repo-type dataset \
    --local-dir metadata \
    --include "*.json"

This populates metadata/grpo_sharerobot_trajectory_{train,test}.json.

mkdir pretrained_weights && cd pretrained_weights
git clone https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct
cd ..

The runner expects a machine with N GPUs; the paper used 4×H100. The script auto-detects N via nvidia-smi.

export THINK_NUM=4       # K thoughts per prompt
export ANSWER_NUM=4      # M post-thought continuations per thought
bash scripts/run_grpo_lora.sh

With K=M=4 on 4 GPUs, each GPU draws K=4 thoughts and K*M=16 continuations from the same prompt per step (the synchronized sampler hands the same prompt to every rank), so one optimizer step processes 4×K=16 thoughts and 4×K×M=64 continuations globally for that prompt. See scripts/README.md for every knob exposed via environment variables.

GRPO-MA/
├── config/                 # GRPOConfig, GRPOScriptArguments, GRPOModelConfig
├── dataset/
│   └── dataset_grpo.py     # YAML-driven multi-image/video GRPO dataset
├── metadata/                  # JSON manifests (downloaded from HF, not in repo)
│   └── README.md              # Where to fetch them and the schema
├── model/
│   ├── qwen_module.py             # Qwen2.5-VL wrapper
│   ├── vlm_module.py              # Abstract VLM interface
│   ├── qwen2_5vl_monkey_patch.py  # Flash-attn / DS-Zero3 / torch.load patches
│   ├── reward_func.py             # Routes (completion, sol, question_type) → task rewards
│   └── task_configs.py            # Reads task templates from the registry
├── scripts/
│   ├── run_grpo_lora.sh           # Training entry point
│   ├── eval/eval.py               # Multi-task evaluation driver
│   ├── eval/run_eval_trajectory.sh
│   ├── smoke_test.py              # No-GPU sanity check
│   ├── train/grpo_trajectory.yaml # Training data manifest
│   └── zero2.json                 # DeepSpeed ZeRO-2 config
├── task/
│   ├── task_loader.py             # Auto-discovers task modules under task/
│   ├── processing_utils.py        # Coordinate-rescaling helpers
│   ├── trajectory_sharerobot.py   # The reference task
│   └── TEMPLATE.py                # Copy this when adding a new task
├── trainer/
│   └── grpo_ma_trainer.py         # GRPOMATrainer (K×M sampling, dual advantages)
├── train.py                       # CLI entry point
└── setup.py

scripts/run_grpo_lora.sh is the canonical entry point. The most important knobs:

Set GRPO_MA_VERBOSE=1 to re-enable per-step diagnostic prints (sampler indices, generated completions, sign-inconsistent counts). They are silenced by default to keep training logs readable. Full hyperparameter documentation is in scripts/README.md.

The training runs reported in the paper used K=M=4, lr=1e-5, LoRA rank 64, frozen vision tower, DeepSpeed ZeRO-2, bf16, FlashAttention-2, and a </think> stop string.

The evaluation driver is task-agnostic. Each task module owns its scoring via a compute_metrics(predictions, **kwargs) callable that is auto-registered by task/task_loader.py.

python scripts/eval/eval.py \
    --model_name_or_path pretrained_weights/Qwen2.5-VL-3B-Instruct \
    --adapter_path output/<run>/checkpoint-<step> \
    --dataset metadata/grpo_sharerobot_trajectory_test.json \
    --image_root data/ShareRobot/trajectory/images \
    --output_dir output/eval/<run>

Or use the convenience wrapper:

ADAPTER_PATH=output/<run>/checkpoint-<step> bash scripts/eval/run_eval_trajectory.sh

Outputs:

output/eval/<run>/
├── predictions.jsonl    # One record per test item with raw + parsed prediction
└── metrics.json         # Aggregated metrics, grouped by task_type

For the trajectory task this reports DFD, Hausdorff distance, RMSE, endpoint distance, and format compliance (matching Table 4 in the paper).

To add evaluation for a new task: implement compute_metrics(...) in the task module — see task/trajectory_sharerobot.py for the reference implementation.

The framework auto-discovers anything you drop into task/, so a new task is just one file plus a metadata manifest.

Place media (images, videos) under data/ and create a metadata JSON/JSONL inside metadata/. Each record needs at minimum:

• question: the prompt shown to the model
• answer: the ground-truth answer (string, number, JSON, etc.)
• question_type: must match TASK_CONFIG["task_type"] in your task module
• Optional media fields: image, video, plus any task-specific attributes

Example (metadata/my_task_train.json):

[
  {"question": "Which tool is highlighted?", "answer": "hammer", "image": "train/toolbox_001.jpg", "question_type": "my-task"},
  {"question": "Locate the object.", "answer": "[120, 48, 256, 220]", "image": "train/toolbox_002.jpg", "question_type": "my-task"}
]

# scripts/train/grpo_my_task.yaml
datasets:
  - json_path: metadata/my_task_train.json
    sampling_strategy: "all"
    data_root: data/my_task/images
    data_modality: image

cp task/TEMPLATE.py task/my_task.py

Edit:

• TASK_CONFIG["task_type"] — unique lowercase-with-hyphens identifier, matching question_type.
• TASK_CONFIG["grpo_template"] — prompt template. {Question} is replaced with the question text. Other config keys (description, input_format, etc.) are logged but do not affect training.
• format_reward(completion, sol, **kwargs) — cheap structural check (returns 0.0 or 1.0).
• accuracy_reward(completion, sol, **kwargs) — real task metric. May return float or (float, dict); the dict is logged as auxiliary metrics during training.
• process_answer(...) (optional) — rescale pixel-space answers when the image is resized (use task.processing_utils.scale_trajectory / scale_bbox_xyxy / scale_points).
• compute_metrics(predictions, **kwargs) (optional) — used by the eval driver.

register(...) is auto-called by the task loader; leave it as-is.

python -c "from task import print_summary; print_summary()"
python scripts/smoke_test.py

export DATASET_NAME=scripts/train/grpo_my_task.yaml
export RUN_NAME=Qwen2.5-VL-3B-GRPO-lora-my-task
bash scripts/run_grpo_lora.sh

• All paper experiments used Qwen2.5-VL-3B-Instruct (also 7B for the scaling study) with LoRA (rank 64, alpha 128, dropout 0.05) on a 4×H100 80GB node.
• We did not use vLLM during training; generation runs through model.generate inside the trainer.
• Random seed: 42 (set via --data_seed and set_seed).
• The warmup_ratio is 0.0 in the paper. The trainer has logic that forces M=1 during warmup (degrading GRPO-MA to GRPO); this is unused with warmup_ratio=0.0.
• The released ShareRobot trajectory metadata uses id, image, question, question_type, answer fields. Images are loaded relative to the YAML data_root.

If you find this work useful, please cite:

@inproceedings{wang2026grpoma,
  title     = {Why Tree-Style Branching Matters for Thought Advantage Estimation in {GRPO}},
  author    = {Wang, Hongcheng and Huang, Yinuo and Wang, Sukai and Ren, Guanghui and Dong, Hao},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  series    = {PMLR},
  year      = {2026},
  url       = {https://arxiv.org/abs/2509.24494}
}

The trainer is built on top of the open-source GRPO implementation in HuggingFace TRL. Vision-language model wrappers and the Flash-Attention monkey patches build on HuggingFace Transformers and PEFT. The DAPO-style clip-higher and several GRPO stabilizers we compare against were proposed by their respective authors; see the paper for full citations. The ShareRobot dataset is released by BAAI.

Released under the Apache License 2.0. Portions derived from third-party Apache-2.0 projects are attributed in NOTICE.