On Training in Imagination

• URL: http://arxiv.org/abs/2605.06732v2
• Date: Mon, 11 May 2026 21:19:23 GMT
• Title: On Training in Imagination
• Abstract summary: State-of-the-art model-based reinforcement learning methods train policies on imagined rollouts.<n>We study how errors in learned dynamics and reward models affect returns and policy optimization.
• Score: 69.97419830683606
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: State-of-the-art model-based reinforcement learning methods train policies on imagined rollouts. These rollouts are trajectories generated by a learned dynamics model and are scored by a learned reward model, but without querying the true environment during policy updates. We study this training paradigm by quantifying how errors in learned dynamics and reward models affect returns and policy optimization. First, we extend the analysis of Asadi et al. (2018) to MDPs with learned reward models, and derive the optimal sample allocation--the ratio of dynamics samples to reward samples that minimizes a bound on return error under power-law scaling assumptions. We identify lower Lipschitz constants of the learned dynamics, reward, and policy as a representation desideratum that tightens this bound, and we connect this perspective to the temporal-straightening objective of Wang et al. (2026). Second, we examine how policy optimization with REINFORCE tolerates noisy rewards, which are often cheaper to obtain. We show that zero-mean reward noise leaves the gradient estimator unbiased and adds at most a variance term that decreases with the number of rollouts. This introduces a practical tradeoff: given a fixed budget, should one buy more rollouts with cheaper but noisier rewards, or fewer rollouts with more expensive but less noisy rewards? We reduce this choice to a one-dimensional optimization problem and characterize the optimum.

Related papers

• Your Language Model is Its Own Critic: Reinforcement Learning with Value Estimation from Actor's Internal States [14.658333957936769]
  We introduce Policy Optimization with Internal State Value Estimation.<n>It obtains a baseline at negligible cost by using the policy model's internal signals already computed during the policy forward pass.<n>A lightweight probe predicts the expected verifiable reward from the hidden states of the prompt and generated trajectory.
  arXiv  Detail & Related papers  (2026-05-08T10:49:36Z)
• MARBLE: Multi-Aspect Reward Balance for Diffusion RL [71.6241143519038]
  Reinforcement learning has become the dominant approach for aligning diffusion models with human preferences.<n>Existing practice deal with multiple rewards by training one specialist model per reward.<n>We propose MARBLE, a gradient-space optimization framework that maintains independent advantage estimators for each reward.
  arXiv  Detail & Related papers  (2026-05-07T16:20:42Z)
• Real-Time Aligned Reward Model beyond Semantics [49.717236911878224]
  We introduce R2M (Real-Time Aligned Reward Model), a novel lightweight RLHF framework.<n>R2M goes beyond vanilla reward models that solely depend on the semantic representations of a pretrained LLM.<n>This work points to a promising new direction for improving the performance of reward models through real-time utilization of feedback from policy models.
  arXiv  Detail & Related papers  (2026-01-30T07:32:35Z)
• GDRO: Group-level Reward Post-training Suitable for Diffusion Models [55.948229011478304]
  Group-level rewards successfully align the model with the targeted reward.<n>Group-level Direct Reward Optimization (GDRO) is a new post-training paradigm for group-level reward alignment.<n>GDRO supports full offline training that saves the large time cost for image rollout sampling.<n>It is diffusion-sampler-independent, which eliminates the need for the ODE-to-SDE approximation to obtainity.
  arXiv  Detail & Related papers  (2026-01-05T11:47:18Z)
• Nested-ReFT: Efficient Reinforcement Learning for Large Language Model Fine-Tuning via Off-Policy Rollouts [25.205293698698867]
  We introduce Nested-ReFT, where a subset of layers of the target model acts as the behavior model to generate off-policy completions during training.<n>Our theoretical analysis shows that Nested-ReFT yields unbiased gradient estimates with controlled variance.<n>Our empirical analysis demonstrates improved computational efficiency measured as tokens/sec across multiple math reasoning benchmarks and model sizes.
  arXiv  Detail & Related papers  (2025-08-13T18:37:46Z)
• Probabilistic Uncertain Reward Model [27.40414952747553]
  This paper proposes the Probabilistic Uncertain Reward Model (PURM) to learn the reward distributions that emerged from the preference data.<n>We show that PURM outperforms existing methods with more accurate reward and sound uncertainty estimations.
  arXiv  Detail & Related papers  (2025-03-28T14:39:52Z)
• RED: Unleashing Token-Level Rewards from Holistic Feedback via Reward Redistribution [50.171320156632866]
  Reinforcement learning from human feedback offers a promising approach to aligning large language models with human preferences.<n>Current reward models operate as sequence-to-one models, allocating a single, sparse, and delayed reward to an entire output sequence.<n>We propose a more fine-grained, token-level guidance approach for RL training.
  arXiv  Detail & Related papers  (2024-11-13T02:45:21Z)
• Robust Preference Optimization through Reward Model Distillation [68.65844394615702]
  Direct Preference Optimization (DPO) is a popular offline alignment method that trains a policy directly on preference data.<n>We analyze this phenomenon and use distillation to get a better proxy for the true preference distribution over generation pairs.<n>Our results show that distilling from such a family of reward models leads to improved robustness to distribution shift in preference annotations.
  arXiv  Detail & Related papers  (2024-05-29T17:39:48Z)
• Fine-Tuning Language Models with Reward Learning on Policy [68.70065254564642]
  Reinforcement learning from human feedback (RLHF) has emerged as an effective approach to aligning large language models (LLMs) to human preferences. Despite its popularity, (fixed) reward models may suffer from inaccurate off-distribution. We propose reward learning on policy (RLP), an unsupervised framework that refines a reward model using policy samples to keep it on-distribution.
  arXiv  Detail & Related papers  (2024-03-28T10:02:10Z)
• Scaling Laws for Reward Model Overoptimization [19.93331579503503]
  We study how the gold reward model score changes as we optimize against the proxy reward model using either reinforcement learning or best-of-$n$ sampling. We also study the effect on this relationship of the size of the reward model dataset, the number of reward model and policy parameters, and the coefficient of the KL penalty added to the reward in the reinforcement learning setup.
  arXiv  Detail & Related papers  (2022-10-19T17:56:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.