Think Outside the Policy: In-Context Steered Policy Optimization

AI summary

Overview

Research area: Reinforcement learning for large reasoning models (RLVR — Reinforcement Learning from Verifiable Rewards), specifically reward-based post-training of LLMs for mathematical reasoning.

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods (GRPO, PPO), importance sampling, off-policy corrections, and in-context learning theory.

Scope: The paper proposes In-Context Steered Policy Optimization (ICPO), a framework that uses a model's own in-context learning ability to generate "expert-guided" rollouts from existing datasets, replacing the need for stronger external models to expand exploration during RLVR training.

What This Paper Is About

Standard RLVR methods such as Group Relative Policy Optimization (GRPO) only learn from on-policy rollouts, so exploration stays trapped inside the current policy's distribution and trajectory diversity is narrow. Recent fixes import trajectories from stronger "expert" models, but this is expensive and those models are often unavailable. ICPO instead treats in-context demonstrations drawn from existing datasets as an implicit expert signal, producing off-policy-style rollouts from the same policy that steer it outside its default distribution.

Key Contributions

1. Empirical validation of ICL rollouts as expert signals. The authors show that 1-shot in-context-conditioned rollouts are more accurate, more diverse (larger inter-trajectory edit distances), and more likely to "flip" previously incorrect answers to correct ones than temperature-based sampling alone (temperature raised from 0.6 to 1.2).

2. Mixed-policy GRPO with Implicit Expert Forcing (IEF). GRPO is extended so each rollout group mixes on-policy trajectories with expert-conditioned trajectories generated by the same policy under an in-context demonstration prefix. All rollouts come from one model, so this is described as an "input-conditioned off-policy" method rather than a model-based off-policy one.

3. Expert Region Reject Sampling (ERRS). A filtering step that keeps only expert-conditioned trajectories whose reward exceeds a threshold δ (set to 1.0 by default), preventing noisy off-policy traces from contaminating policy updates.

4. Reward Shaping with an annealed expert bonus (ICPO†). A step-annealed bonus α·γ(t), where γ(t) = 1 − t/T decays linearly over training and α is set to 1.0, applied only to correct trajectories inside the expert region. This provides strong early expert guidance that relaxes over time.

Main Findings

• In-distribution gains over GRPO. On Qwen3-1.7B, ICPO reaches an in-distribution average of 52.5 versus 48.4 for GRPO (+4.1); ICPO† reaches 51.4 (+3.0). On Qwen3-8B, ICPO scores 65.7 versus 63.5 (+2.2) and ICPO† scores 65.0 (+1.5). The abstract reports maximum average improvements of up to +4.1 and +4.0 points over vanilla GRPO and mixed-policy GRPO respectively.

• Out-of-distribution trade-off. Adding reward shaping gives ICPO† a more balanced profile: +0.7 OOD average on Qwen3-1.7B and +2.4 on Qwen3-8B, whereas plain ICPO on Qwen3-1.7B is -1.9 OOD despite the largest in-distribution gain.

• Comparison to prior methods. On Qwen2.5-Math-7B, ICPO† attains a 53.4 average and ICPO 53.2, versus ReLIFT (48.8), LUFFY (50.1), and Prefix-RFT (51.8) — reported as gains of +4.6, +3.3, and +1.6 points respectively.

• More rollouts are not enough. GRPO with 16 on-policy rollouts (GRPO_ExtraRollouts) yields only +2.3 on Qwen3-1.7B and +0.8 on Qwen3-8B in-distribution, with -0.8 and +1.0 OOD, showing that brute-force rollout amplification does not change the underlying policy distribution.

• Component ablations (Qwen3-8B). Removing IEF costs 1.2 average points and removing ERRS costs 0.8, while removing both restores plain GRPO performance — every component contributes.

• Robustness to expert data choice. Replacing in-domain CoT math demonstrations with cross-domain Program-of-Thought (PoT) code data still beats GRPO: 51.5 versus 48.4 on Qwen3-1.7B, and 65.1 versus 63.8 on Qwen3-8B.

• No dependence on pre-existing ICL skill. On Qwen3-8B-Base, described as a model with no ICL capability, ICPO raises Pass@1 from 48.8 (GRPO) to 50.5 and Pass@32 from 68.3 to 74.9, suggesting ICPO expands the exploration space rather than sharpening existing modes.

• Distribution shift is measurable. Instance-level perplexity analysis on MATH-500 shows ICPO points fall below the y = x diagonal relative to GRPO, and token-rank shift analysis shows GRPO collapsing back to high-probability base-policy outputs beyond the 60th percentile while ICPO decays more slowly, with ICPO† showing the largest late-generation shift.

• Higher entropy during training. ICPO maintains higher policy entropy than GRPO, reflecting broader policy support.

• Compute cost. Training took roughly 3 days for the 1.5B-scale model and 7 days for the 8B model on 8 NVIDIA H20 GPUs.

Methodology in Plain English

The setup works like this. For each training prompt, the model normally samples 8 rollouts and GRPO normalizes rewards within that group to compute advantages. ICPO changes the group composition: it keeps 7 on-policy rollouts and generates 1 additional rollout by prepending k demonstration question-answer pairs sampled from the MATH training set (7.5k problems with high-quality solutions) to the prompt. Because the input is different, this rollout is treated as off-policy and gets its own importance weight relative to the expert-conditioned distribution.

The authors justify this with a hypothesis-class view of in-context learning: a Transformer applied to [demonstrations; query] decomposes into a task-vector function A(D) and a reasoning function F(q; ϑ), so the demonstrations implicitly inject a task vector ϑ that bends the output distribution toward expert-like behavior without any parameter update.

Two stabilizers are layered on top. First, reject sampling: an expert-conditioned rollout is only substituted into the group if its verifiable reward is at least δ = 1.0 (correct answers), so bad off-policy traces never contribute gradients. Second, an optional reward bonus (ICPO†) that adds α·γ(t) to correct expert-region rollouts, decaying linearly to zero over training so early imitation gives way to autonomous optimization. Losses use the standard PPO-style clipping function with a shaping transform f(x) = x/(x + λ) where λ = 0.01.

Training used the filtered OpenR1-Math-220k subset (approximately 45k verified reasoning prompts, excluding generations over 8192 tokens and those Math-Verify marked incorrect). Evaluation covered AIME24, AIME25, AMC23, Minerva, Olympiad, and MATH-500, with Avg@32 for the small AIME and AMC sets and Pass@1 elsewhere, plus out-of-distribution tests on ARC-C, GPQA-Diamond, and MMLU-Pro with shuffled multiple-choice options.

Why This Matters

Impact on research. The paper reframes in-context learning as a training-time exploration mechanism rather than an inference-time trick. It also offers a cheaper alternative axis for expanding exploration in RLVR — input conditioning instead of external model rollouts — and provides distributional-shift diagnostics (per-instance perplexity, token-rank shift) that other RLVR papers could adopt.

Real-world applications:

• Post-training of open-weight reasoning models where no larger teacher model is available or affordable.
• Mathematical tutoring and step-by-step solution systems that need broader solution-space coverage.
• Cross-domain steering, since the PoT experiment suggests demonstrations from one domain (code) can reshape behavior in another.
• Low-cost RLVR pipelines for teams without large GPU clusters, because guidance comes from datasets rather than from generating traces with a frontier model.

Industry relevance. The method is directly compatible with existing GRPO infrastructure and drops into standard training libraries; the reported code release (github.com/Celine-hxy/ICPO) lowers adoption cost. The stated hardware footprint (8 H20 GPUs, 3–7 days) is within reach of many industrial labs.

Future Directions

• Extending ICPO beyond mathematical and knowledge reasoning to domains such as code generation, which the authors flag as likely requiring task-specific reward design and evaluation.
• Clarifying the mechanism by which ICPO helps models without ICL ability, since the Qwen3-8B-Base results suggest improvement even when the in-context mechanism should not apply — the authors describe ICL as a "scaling amplifier" rather than a prerequisite.
• Quantifying the trade-off between in-distribution gains and OOD losses, since plain ICPO's largest in-distribution win on Qwen3-1.7B came with a -1.9 OOD average.
• Tuning the design knobs left fixed here — the reject threshold δ = 1.0, bonus weight α = 1.0, shaping λ = 0.01, the 1-off-policy-to-7-on-policy ratio, and the number of demonstrations k — is left largely unexplored.

Target Audience

Researchers and engineers working on RL post-training of large reasoning models, particularly those implementing GRPO-style pipelines and looking for cheaper ways to expand exploration. It is also relevant to readers interested in the mechanistic role of in-context learning during training, and to practitioners who want a plug-and-play, dataset-driven alternative to distilling trajectories from stronger models. Readers without a background in policy-gradient RL will find the method sections difficult.

Authors’ abstract

Existing Reinforcement Learning from Verifiable Rewards (RLVR) methods, such as Group Relative Policy Optimization (GRPO), have achieved remarkable progress in improving the reasoning capabilities of Large Reasoning Models (LRMs). However, they exhibit limited exploration due to reliance on on-policy rollouts which are confined to the current policy's distribution, resulting in narrow trajectory diversity. Recent approaches attempt to expand policy coverage by incorporating trajectories generated from stronger expert models, yet this reliance increases computational cost and such advanced models are often inaccessible. To address these issues, we propose In-Context Steered Policy Optimization (ICPO), a unified framework that leverages the inherent in-context learning capability of LRMs to provide expert guidance using existing datasets. ICPO introduces mixed-policy GRPO with implicit expert forcing, which expands exploration beyond the current policy distribution without requiring advanced LRM trajectories. To further stabilize optimization, ICPO integrates expert region reject sampling to filter unreliable off-policy trajectories and annealed expert-bonus reward shaping to balance early expert guidance with later autonomous improvement. Results demonstrate that ICPO consistently enhances RLVR performance and training stability on mathematical reasoning benchmarks, revealing a scalable and effective RLVR paradigm for LRMs. Our code is available at https://github.com/Celine-hxy/ICPO.