GRPO-Style Online RL Pipeline
- The paper introduces a pipeline that leverages group relative policy optimization to enable stable, on-policy RL fine-tuning across modular LM programs.
- It groups LM calls by module and invocation index to ensure consistent gradient updates and effective variance reduction through normalized, group-relative advantage scaling.
- The integration of prompt optimization with RL fine-tuning boosts accuracy by up to 11%, demonstrating its practical value in handling complex, variable-length trajectories.
A GRPO-style Online RL pipeline is an on-policy reinforcement learning framework leveraging Group Relative Policy Optimization (GRPO) for credit assignment and stable post-training, originally proposed for LLMs and extended to modular LLM (LM) programs. The multi-module GRPO (mmGRPO) generalizes the standard algorithm to handle structured LM systems composed of multiple modules, each parameterized by distinct prompt templates and potentially separate LoRA adapter weights. The pipeline yields stable, credit-aware RL fine-tuning across modular and heterogeneous LM programs by grouping LM calls for gradient updates and integrating prompt optimization as a staged curriculum (Ziems et al., 6 Aug 2025).
1. Modular LM Program Architecture and Trajectory Semantics
A modular LM program consists of a set of Python "modules" , each equipped with:
- Prompt template (e.g., chain-of-thought, few-shot).
- Module-specific LM weights (commonly sharing a base LM with independent LoRA adapters).
For a structured input , executing yields a final output and trajectory , with each recording the invoked module, the realized prompt (obtained by instantiating 0), and the sampled LM continuation 1 from 2. The system is designed to support complex LM programs with variable-length, interrupted, or modularly compositional execution traces. During training, 3 rollouts 4 are collected, optionally mixing student and teacher programs.
The core idea is to align module calls across rollouts by 5—identifying, for instance, "the third call to module QueryGenerator"—such that all 6's with a shared 7 index are grouped for RL optimization.
2. mmGRPO Policy Objective and Credit Assignment
For each module 8 and invocation index 9, the mmGRPO pipeline forms a group 0, where 1 is the group size. Each tuple receives a program-level reward 2 (where 3 is a chosen metric, such as accuracy). The module parameters 4 are optimized via a clipped, group-relative policy gradient:
5
where:
- 6 is the token-level importance ratio.
- The group-relative advantage is
7
- 8 is the PPO-style clipping parameter (default 9).
- 0 controls KL-regularization to a reference policy.
This objective replaces value function estimates with directly normalized within-group advantage scaling, yielding variance reduction and robust credit assignment even under variable trajectory lengths or changing group composition.
3. Handling Variable-Length and Interrupted Trajectories
Real-world modular LM programs produce trajectories where the number of calls per module and their positions may vary due to early termination or errors. To address this:
- Trajectories are bucketed by 1, producing groups of potentially uneven size.
- Two padding modes are supported:
- "truncate": drop 2 groups whose count is below 3
- "fill": pad missing entries by duplicating or inserting dummy entries (with a low fallback reward)
- Each group is resampled up/down to a fixed 4 using a diversity-promoting sampler (e.g., SelectK-DiverseElements), favoring broad reward coverage and mitigating overfitting to trivial batch statistics.
4. Prompt Optimization and Curriculum Integration
The mmGRPO pipeline may be preceded by integrated prompt optimization to synthesize improved prompt templates 5. This is typically done using a black-box optimizer (e.g., MIPROv2). The staged protocol is:
- Run prompt optimization: 6.
- Fix prompt templates and conduct RL fine-tuning: 7.
Empirically, this BetterTogether(PO, mmGRPO) approach yields stronger initial rollouts and boosts final accuracy across multiple task types, including many-hop search and privacy-preserving delegation (Ziems et al., 6 Aug 2025).
5. End-to-End Pipeline and Pseudocode
The mmGRPO training loop consists of the following steps:
- Sample a batch of 8 inputs from the training dataset.
- For each input, generate 9 rollouts (possibly including multiple teacher programs) to build the experience set.
- Align and group all module invocations across rollouts into 0-indexed groups.
- Pad or truncate groups to the target size 1 and sample for diversity.
- For each group 2, compute group-relative advantages 3 and take one gradient step on 4 using the mmGRPO objective.
- Repeat for 5 training iterations.
The function FormModuleLevelGroups constructs 6 groups from trajectory logs, pads/truncates, and applies diversity sampling as required. Trained models can be deployed or evaluated immediately after this loop completes.
6. Key Hyperparameters and Practical Tips
Crucial settings for stable, high-performance mmGRPO-style online RL include:
| Parameter | Recommended Value | Effect/Notes |
|---|---|---|
| Rollouts per example 7 | 12 | Matching group size 8 for variance control |
| Group size 9 | 12 | Consistency across modules/groups |
| Training steps | 0 | |
| Batch size 1 | 4 | 2 inputs 3 rollouts = 4 trajectories/step |
| Learning rate 5 | 6 | |
| Gradient accum. | 20 | Imitates batch size 20 |
| Weight decay 7 | 0 | |
| Clipping 8 | 0.2 (PPO-style) | Trust region enforcement |
| KL penalty 9 | 0.01–0.04 | Proportional to model size/task |
| Context cutoff | 8192 tokens | Drop trajectories exceeding this length |
| LoRA adapters | rank=16, alpha=64 | Layerwise attention projection tuning |
| Sampling temp. | 0.6 | Exploration/exploitation balance |
| Top-p, top-k | 0.9–0.95, 20 | Search diversity control |
Additional strategies include setting pad_mode="fill" to retain group diversity and using group diversity sampling (reward stratification or variance maximization) to prevent collapse.
7. Implementation, Extensibility, and Empirical Impact
mmGRPO is implemented in the DSPy framework under the optimizer dspy.GRPO. Extensibility is supported via pluggable grouping strategies, padding modes, diversity samplers, and rollout teacher mixes. DSPy’s Arbor training library (GRPOTrainer) manages gradient accumulation, KL logging, and checkpointing. The modular APIs allow integration with custom programs built from DSPy’s module and chain-of-thought primitives.
Empirically, mmGRPO composed with automatic prompt optimization yields substantial improvements: up to 11% accuracy gains over post-trained LMs, and +5% improvement versus prompt optimization alone, across diverse tasks. The method stably supports arbitrary modular LM architectures and has become a reference pipeline for robust, credit-aware online RL with complex programmatic LMs (Ziems et al., 6 Aug 2025).