# On Training in Imagination

arXiv: https://arxiv.org/abs/2605.06732 · HTML: https://mlsanity.com/paper/2605.06732/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 25 of its 41 references that are in the map.
- First posted on arXiv 7 May 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)
- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 25 of 25)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)
- [Reward-Free Exploration for Reinforcement Learning](https://mlsanity.com/paper/2002.02794/) (2020)
- [Training Agents Inside of Scalable World Models](https://mlsanity.com/paper/2509.24527/) (2025)
- [Lipschitz Continuity in Model-based Reinforcement Learning](https://mlsanity.com/paper/1804.07193/) (2018)
- [Reinforcement Learning with Verifiable yet Noisy Rewards under Imperfect Verifiers](https://mlsanity.com/paper/2510.00915/) (2025)
- [Scaling laws for single-agent reinforcement learning](https://mlsanity.com/paper/2301.13442/) (2023)
- [Recurrent Orthogonal Networks and Long-Memory Tasks](https://mlsanity.com/paper/1602.06662/) (2016)
- [Temporal Straightening for Latent Planning](https://mlsanity.com/paper/2603.12231/) (2026)
- [Scaling Laws for Pre-training Agents and World Models](https://mlsanity.com/paper/2411.04434/) (2024)
- [Active Measure Reinforcement Learning for Observation Cost Minimization](https://mlsanity.com/paper/2005.12697/) (2020)
- [An Optimal Tightness Bound for the Simulation Lemma](https://mlsanity.com/paper/2406.16249/) (2024)
- [Learning the Reward Function for a Misspecified Model](https://mlsanity.com/paper/1801.09624/) (2018)
- [Non-Uniform Noise-to-Signal Ratio in the REINFORCE Policy-Gradient Estimator](https://mlsanity.com/paper/2602.01460/) (2026)

## Abstract (arXiv)

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.