On Training in Imagination

Authors: Nadav Timor, Ravid Shwartz-Ziv, Micah Goldblum et al.

arXiv · PDF

Summary

The authors study model-based RL, where policies are trained on imagined rollouts from learned dynamics and reward models without querying the real environment. They analyze how errors in the learned models affect returns and derive the optimal sample allocation—how many samples to spend learning dynamics vs. rewards—under power-law scaling. They also show that zero-mean reward noise doesn't bias the REINFORCE gradient, only adds variance, which creates a tradeoff: is it better to get more rollouts with cheap noisy rewards, or fewer rollouts with expensive accurate rewards?

Main takeaways:

• Extends error-propagation analysis to MDPs with learned reward models and derives optimal dynamics-vs-reward sample allocation
• Identifies lower Lipschitz constants (smoother dynamics, reward, and policy) as a representation goal that tightens error bounds
• Shows zero-mean reward noise leaves REINFORCE gradient unbiased, only adding variance that decreases with more rollouts
• Frames the noisy-reward tradeoff as a one-dimensional optimization: more cheap noisy rollouts vs. fewer expensive accurate ones

Relevance

Not directly related to my persona or midtraining work—this is about model-based RL for sequential decision-making. Could be loosely relevant if I ever modeled behavior installation as a learned dynamics problem (e.g., "how does the model's output evolve under a persona prompt?"), but that's a stretch. Mostly included for general RL background.

Abstract

arXiv:2605.06732v1 Announce Type: new 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.