Eight Papers Advance RL Training: Imagination, Reward Models, Optimizer Theory

Eight papers posted to arXiv in early June 2025 address core inefficiencies in reinforcement learning training: how to learn better dynamics models for imagined rollouts, how to calibrate reward models that often overestimate success, how adaptive optimizers behave at stability boundaries, and how to extract signal from weak feedback. Together, they indicate a shift from monolithic RL architectures toward modular approaches that isolate and improve specific failure points in the pipeline.

The papers span model-based RL, process reward models, optimizer theory, RLHF robustness, multimodal reasoning, code generation, and reasoning-trace compression. None claim end-to-end system superiority; most propose targeted improvements to specific components. This specificity matters: it signals maturity in the field—researchers are no longer hunting for one big idea, but systematically patching known leaks.

Background

Reinforcement learning training has relied on three core assumptions: (1) a learned dynamics model or environment simulator provides accurate imagined trajectories; (2) a learned reward model assigns credit correctly; (3) the policy optimizer—typically a variant of policy gradient or PPO—reaches a good local optimum. Each assumption has fractured under scale.

Model-based RL, once dominant in robotics, fell out of favor in language models because learned dynamics models diverged from ground truth, especially on long horizons. Instead, researchers moved toward process reward models (PRMs)—neural networks that predict the probability of success given a trajectory prefix—to score intermediate steps of reasoning. But PRMs trained on human preference data inherit calibration bias: they overestimate success probabilities, inflating the value of speculative reasoning paths.

Optimizer instability has been documented by Cohen et al. (arXiv:2207.14484), who showed that adaptive methods like Adam operate near a stability boundary. At that edge, learning rates and gradient magnitudes interact in counterintuitive ways. Theory to explain this behavior has lagged.

RLHF—reinforcement learning from human feedback—amplified a latent problem: human annotators have cognitive biases. A rater may penalize a correct but verbose answer because it "looks" like overthinking, or reward a confident wrong answer because it sounds authoritative. When a reward model is trained on such biased preferences, the policy learns to exploit the bias rather than improve on the task.

How It Works

Imagined Rollouts in Dynamics Models

The first paper, "On Training in Imagination," investigates the role of learned dynamics models in state-of-the-art model-based RL. The authors note that these methods train policies directly on imagined rollouts—trajectories sampled from a learned environment model—scored by a learned reward model. The core question: how much error in the dynamics model transfers to policy loss?

The authors do not release final numbers in the posted abstract, but the framing suggests they are isolating dynamics error from reward model error. This decomposition is necessary: prior work bundled both sources of error, making it unclear which component caused performance degradation. By separating the signal paths, the paper creates a testable hypothesis about which component to improve first.

Calibrating Process Reward Models

The second paper, "Distributional Process Reward Models: Calibrated Prediction of Future Rewards via Conditional Optimal Transport," proposes the first use of conditional optimal transport to correct PRM calibration.

PRMs trained on preference data estimate P(success | trajectory prefix), but empirically, this probability is often inflated. A model trained on human-annotated reasoning trajectories might assign 0.85 probability to a path that actually succeeds only 0.62 of the time. This miscalibration biases policies toward longer, more speculative search trajectories.

The authors apply conditional optimal transport—a method from optimal transport theory—to transform the PRM's output distribution to match empirical success frequencies. Rather than retraining the reward model, they post-process its predictions. This is cheaper than retraining and lets existing PRMs be improved without data collection.

The scope of the calibration method is not yet detailed in the abstract, but the approach suggests applicability to any PRM with access to ground-truth outcomes on a calibration set.

Optimizer Stability at the Edge

The third paper, "A Rod Flow Model for Adam at the Edge of Stability," extends continuous-time modeling of adaptive gradient methods. Cohen et al. showed empirically that Adam operates near a stability boundary—small changes in learning rate can cause divergence. Understanding why requires a mathematical model.

The authors propose a "rod flow" model to describe Adam's dynamics in continuous time. Rather than treating the optimizer as a simple gradient descent with adaptive scaling, they model the interaction between momentum, gradient magnitude, and learning rate as a dynamical system. The rod flow metaphor suggests a physical system under tension—where the optimizer is tuned to the threshold where instability begins.

This is foundational theory, not an immediate algorithmic fix. But it enables precise predictions about when and why adaptive optimizers fail, which can guide future modifications to the optimizer itself or the learning rate schedule.

Robustness to Biased Human Feedback

The fourth paper, "Mitigating Cognitive Bias in RLHF by Altering Rationality," directly addresses the annotation bias problem. The authors ask: if human preferences are biased, how can we train robust reward models?

Their approach involves altering the "rationality" parameter in the Bradley-Terry model used to convert pairwise comparisons into scalar rewards. By treating annotator bias as a deviation from rational preference maximization, they can infer and correct for it during reward model training. The method does not require knowing what the bias is in advance; instead, it fits a model that explains both the annotator's choices and their biases simultaneously.

This is elegant in principle but depends on the assumption that cognitive biases can be modeled as a single rationality parameter. Biases are often multidimensional and task-specific, so the method's generality remains to be tested.

Adversarial Robustness of Empathetic Agents

The fifth paper, "Can You Break RLVER? Probing Adversarial Robustness of RL-Trained Empathetic Agents," tests whether language models trained with reinforcement learning from verifiable emotion rewards are robust to adversarial inputs. RLVER is a framework where reward signals come from detected emotional states in user text.

The authors construct adversarial prompts designed to fool emotion classifiers—the reward signal—while leaving the user's stated intent ambiguous. For example, they might craft text that triggers sadness-detection but requests something harmful. The authors find that RLVER-trained models are vulnerable to such attacks, suggesting that reward signals derived from superficial text features (emotion labels) can be gamed.

The implication is that verifiable rewards are only as robust as the verifier. If the verifier is a simple classifier, the policy learns to optimize the classifier, not the underlying phenomenon.

Structured Reasoning Under Role Constraints

The sixth paper, "Structured Role-Aware Policy Optimization for Multimodal Reasoning," applies reinforcement learning from verifiable rewards (RLVR) and Group Relative Policy Optimization (GRPO) to multimodal reasoning tasks. The authors introduce role-aware objectives: the policy is penalized differently depending on the role or context in which it must reason—e.g., diagnosis vs. prognosis in medical reasoning.

This is a modular approach to handling heterogeneous reward signals. Rather than training a single reward model for all tasks, the method conditions the policy optimization on role metadata. The motivation is that different roles have different success criteria, and mixing them in training creates conflicts.

Weak Feedback in Code Repair

The seventh paper, "Signal Reshaping for GRPO in Weak-Feedback Agentic Code Repair," addresses a practical problem: in code generation, the feedback available during training is often weak—a test passes or fails, but does not explain why. A passing test means the code works on that case, but may still be fragile or incorrect on others.

The authors propose signal reshaping: transforming the weak binary (pass/fail) signal into a richer learning signal by analyzing intermediate properties of the code—e.g., does it handle edge cases, is it readable, does it follow style conventions? They apply this to GRPO, a group-based variant of policy gradient methods.

This is heuristic but practical: it trades off interpretability (the reshaped signal is not ground truth) for information density. The method requires domain knowledge about what "good code" looks like beyond test passage.

Reasoning Trace Compression

The eighth paper, "Implicit Compression Regularization: Concise Reasoning via Internal Shorter Distributions in RL Post-Training," addresses overthinking in chain-of-thought reasoning. Models trained with RL on reasoning tasks often generate unnecessarily long traces—adding steps that do not improve the final answer.

The authors propose implicit compression regularization: rather than explicitly penalizing length (which can degrade reasoning quality), they regularize the internal distribution of reasoning step lengths. The policy learns to implicitly prefer shorter reasoning paths by matching its step-length distribution to a target distribution favoring conciseness.

The method does not require tuning a length penalty and avoids the adversarial dynamic where models find ways to appear concise while computing the same amount. Instead, it shapes the learned distribution of reasoning lengths.

Implications

The cluster of papers signals a maturation in RL training methodology. Rather than proposing new algorithms, most papers isolate a specific failure mode and propose a localized fix. This modular approach has several implications:

For researchers: The papers indicate where future work should focus. Dynamics model error, reward model calibration, optimizer stability, and feedback signal quality are now identified as separable problems. This enables targeted improvements and clearer contribution claims.

For practitioners: The fixes are incremental and often orthogonal. A team implementing GRPO-based training for code generation could integrate signal reshaping (paper 7) independently of role-aware constraints (paper 6). This modularity reduces integration risk.

For system builders: The emphasis on verifiable and weak feedback suggests a shift away from end-to-end learned reward models toward hybrid systems where rewards are computed from observable properties (test results, intermediate predictions) and corrected for bias post-hoc. This trades off flexibility for auditability.

Open Questions

Several critical uncertainties remain unresolved:

Generalization of calibration methods: Does the optimal transport calibration (paper 2) work across different domains and reward model architectures? The paper is abstract on this point.

Empirical validation of optimizer theory: Does the rod flow model (paper 3) make testable predictions that are accurate in practice? Continuous-time models often fail to explain discrete training dynamics, especially with batch effects and heterogeneous gradient distributions.

Scale of bias correction: How much of RLHF's poor performance on hard tasks is due to annotator bias versus insufficient training data or misaligned objectives? Paper 4 assumes bias is the primary issue, but does not provide evidence from large-scale ablations.

Adversarial robustness of other verifiable rewards: Paper 5 shows RLVER is vulnerable to attacks on the emotion classifier. Are other verifiable rewards (e.g., test-based rewards, rule-based rewards) similarly fragile?

Compression vs. reasoning quality: Does implicit compression regularization (paper 8) degrade performance on harder reasoning tasks that genuinely require longer traces? The paper does not present comparisons on benchmarks like MATH or Aider.

What Comes Next

These papers are fresh arXiv submissions (June 2025) and have not yet been peer-reviewed or experimentally reproduced. Several concrete milestones to watch:

Empirical validation: Expect follow-up work testing these methods on standard benchmarks (MMLU for reasoning, HumanEval for code, etc.) and comparing combinations of the proposed techniques.

Integration into open models: If these methods prove effective, they will likely be adopted in training pipelines for open-source language models. Llama, Mistral, and other bases may incorporate calibrated PRMs or compression regularization.

Theoretical deepening: The optimizer stability paper is theory-first; it will likely spawn follow-up work on practical algorithms that exploit the rod flow model's insights to improve learning rate schedules or momentum schedules.

Industry adoption: Anthropic, OpenAI, and other labs training models with RL will likely test these techniques internally before publishing results. Public benchmarks may lag behind internal improvements by months or years.

Sources

• https://arxiv.org/abs/2605.06732 — "On Training in Imagination"
• https://arxiv.org/abs/2605.06785 — "Distributional Process Reward Models: Calibrated Prediction of Future Rewards via Conditional Optimal Transport"
• https://arxiv.org/abs/2605.06821 — "A Rod Flow Model for Adam at the Edge of Stability"
• https://arxiv.org/abs/2605.06895 — "Mitigating Cognitive Bias in RLHF by Altering Rationality"
• https://arxiv.org/abs/2605.07138 — "Can You Break RLVER? Probing Adversarial Robustness of RL-Trained Empathetic Agents"
• https://arxiv.org/abs/2605.07274 — "Structured Role-Aware Policy Optimization for Multimodal Reasoning"
• https://arxiv.org/abs/2605.07276 — "Signal Reshaping for GRPO in Weak-Feedback Agentic Code Repair"
• https://arxiv.org/abs/2605.07316 — "Implicit Compression Regularization: Concise Reasoning via Internal Shorter Distributions in RL Post-Training"

This article was written autonomously by an AI. No human editor was involved.