Imagination helps only while the imagined data still resembles a world the policy will actually visit.
A Budget-Conscious MPC Loop
This section assumes familiarity with one-step world model prediction covered in section 36.3, particularly the rollout-horizon caution introduced there. The imagination rollout technique developed here is extended in section 37.5, which examines when synthetic data accelerates learning and when it causes value collapse (a failure mode where the learner's estimated returns become systematically too optimistic because they are computed from imagined states the real world never confirms). The latent imagination variant of this idea recurs in Part IX alongside generative world models and structured prediction.
A robot arm trained on two hours of real contact data can, with imagination rollouts, behave as if it practiced for twenty. The trick: branch short synthetic trajectories from states the world model already trusts, then blend those imagined transitions into policy training. This is the technique that made sample-efficient model-based RL practical on physical hardware, where every real trial costs time and wear. Here you will build that branching loop from scratch, choose the horizon that keeps imagination honest, and see exactly how synthetic data composition shifts the learner's value estimates before a single extra real step is taken.
Synthetic data is useful only while it stays tethered to states the model understands. Horizon control is what keeps imagination from turning into dataset corruption.
Short Rollouts, Big Consequences
What if a robot could squeeze twenty hours of practice out of two, simply by letting its world model daydream a few steps past each real experience? That is exactly the bargain Model-Based Policy Optimization (MBPO)-style learning strikes, where a learned dynamics model predicts how the world would respond to an action without the robot actually taking it: real states from the replay buffer (the stored set of past real transitions the agent has already collected) seed short model-generated rollouts. Those imagined transitions augment policy learning while limiting compounding error (the way a small per-step prediction mistake adds to the next step's input, so the mistake grows larger with every additional imagined step rather than staying fixed). The core trade-off is direct. More imagined data can accelerate learning, but only if rollout length stays inside the model's trusted region, the envelope of state space where prediction error remains small enough to trust. Figure 37.4A captures this principle: imagined trajectories branch from real states and stay short enough that the model remains locally credible. Consider the scale this unlocks. On HalfCheetah, the original MBPO experiments reached in roughly 300 real environment steps a performance level that a pure model-free baseline needed 50,000 steps to reach. On a physical robot, that difference turns many hours of wear into a few minutes of it.
A policy trained on imagined states beyond the model's trusted region is not a policy trained on data: it is a policy trained on a rumor. On a physical robot, the trusted region has direct mechanical stakes. If the policy trains on imagined states outside it, it may command motions calibrated to a fantasy dynamics that never existed on real hardware: joints that accelerate faster than the servo can follow, contacts assumed frictionless that are not, or torques that the controller clips silently. The result is not just degraded reward but unsafe execution that can damage the robot or its environment.
The trusted region just described is not fixed in advance; it has to be measured from data, which raises the practical question of how a boundary drawn in continuous state space is actually decided. The model determines this boundary empirically: during real experience collection, it tracks per-state prediction residuals (the difference between the model's predicted next state and the next state actually observed) and marks a state as trusted when held-out next-state error falls below a user-set threshold (commonly 5 percent of the observation range). Imagined rollouts branching from states inside that boundary stay locally coherent because nearby transitions were seen during training; rollouts that wander outside the boundary accumulate residual error multiplicatively with each step, which is why horizon length (the number of imagined steps taken forward from a real seed state before the rollout stops, written \(h\) below) is the primary control. Figure 37.4B contrasts these two cases: a short rollout that stays inside the trusted region against a long one whose per-step error compounds until it drifts into states the real system can never occupy.
One useful mental model is
$$ \mathcal{D}_{\text{train}} = \mathcal{D}_{\text{real}} \cup \mathcal{D}_{\text{model}}^{(h)}, $$
where the imagination horizon \(h\) is deliberately small. This keeps model-generated states near the support of real experience.
Small \(h\) is not caution for its own sake; it is what keeps \(\mathcal{D}_{\text{model}}^{(h)}\) inside the support of \(\mathcal{D}_{\text{real}}\).
That support argument is the mechanism, not a stylistic preference. The learner may recycle real states into nearby imagined futures because the model has seen enough neighboring transitions to stay coherent there. Once synthetic states seed further synthetic states, the training set drifts into regions no real interaction ever grounded, and value estimates turn systematically optimistic.
Branching from real buffer states is a bias-control trick. It keeps the synthetic rollout close to regions where the model has at least some evidence.
Worked Probe
If branching keeps synthetic data honest, the next question is how much of that synthetic data a single horizon choice actually creates. The probe below logs how many synthetic transitions are produced from a replay batch under different imagination horizons. It shows why horizon choice changes dataset composition so quickly.
# Count imagined transitions produced from one replay batch.
replay_batch = 128
horizons = [1, 3, 5]
imagined = {h: replay_batch * h for h in horizons}
ratio_to_real = {h: round(imagined[h] / replay_batch, 1) for h in horizons}
print({"imagined_transitions": imagined, "ratio_to_real": ratio_to_real}){'imagined_transitions': {1: 128, 3: 384, 5: 640}, 'ratio_to_real': {1: 1.0, 3: 3.0, 5: 5.0}}
Read the imagined-transitions counts and ratios as a dataset-composition signal: at horizon 1 the synthetic set exactly matches the real batch, but at horizon 5 it is five times larger. That ratio tells you how much weight model-generated data already carries in training before any explicit mixing ratio is set, which is why horizon is not a cosmetic hyperparameter but a direct control on how much model bias enters the learner.
When you implement imagination rollouts, log the real-to-model transition ratio, the rollout branching source, and the maximum horizon. These three numbers explain a large fraction of success or failure in practice. mbrl-lib and Dreamer-style codebases are useful references because they make the replay-to-imagination contract visible rather than hiding it inside one giant trainer.
Common Failure Modes
Logging those ratios tells you how much synthetic data you have, but not where it goes wrong, so it helps to know the specific ways a well-instrumented pipeline still breaks. A short-horizon imagination pipeline can still fail in three physically grounded ways. First, the dynamics model is least accurate in the contact-rich subspace that matters most for reward. A Franka Panda grasping model trained on free-space reaching typically shows low held-out MSE (mean squared error, the average of squared differences between predicted and observed next states) overall, but its error commonly spikes at the moment of fingertip contact. Imagined rollouts seeded near grasp candidates therefore tend, in practice, to overestimate grasp success. (this pattern is not universal: an ensemble that includes contact-rich training data can narrow the gap considerably) Second, the policy can overfit to imagined states that are geometrically plausible under the model but kinematically unreachable on hardware. A 7-DOF arm model trained on 50 Hz joint-position data may silently generate next states that violate the robot's velocity limits (for a Panda, 2.175 rad/s per joint). The policy then commands motions the real controller will saturate or reject. Third, the training loop can let synthetic transitions dominate the replay mixture without triggering any loss alarm. Model MSE on real held-out data stays low while the synthetic buffer grows. On a Spot quadruped this pattern has typically produced a gait that looked stable in simulation but failed on first contact with real terrain, because imagined foot-contact normals did not replicate the terrain's actual friction distribution.
Checkpoint
So far: short-horizon imagination can still fail in three physically grounded ways: the dynamics model is typically least accurate exactly where contact happens, the policy can exploit imagined states that are geometrically plausible but kinematically unreachable on real hardware, and synthetic transitions can quietly dominate the replay mixture without any loss metric flagging it.
The fix is not to abandon imagination, but to instrument it with hardware-aware checks. Log the source joint configuration for each imagined rollout, the per-joint velocity magnitude of the imagined next state against the robot's rated limits, the fraction of imagined transitions that involve contact (end-effector force above a 2 N threshold, or foot-contact flags for legged systems), and at least one replayed real trajectory where the imagined branch diverged from the sensor stream. Those four signals turn a vague trust problem into a concrete mechanical audit.
When using mbrl-lib's ModelEnv.step() with a custom dynamics model, always call .contiguous() on your observation tensor before passing it in. The function accepts non-contiguous tensors without raising an error but silently produces incorrect next-state predictions, making the resulting imagined rollouts look plausible while being numerically wrong. This failure is especially hard to detect because model loss on held-out real data stays low; only the synthetic transitions are corrupted. Adding a one-line assertion assert obs.is_contiguous() at the start of your rollout loop catches it immediately.
Start at \(h = 1\) and increase by one step only when the model's one-step held-out prediction error (MSE on next-state observations) is below a threshold you set before training, for example 5% of the observation range. In the original MBPO experiments on HalfCheetah (Janner et al., 2019), a horizon of 1 to 3 steps already captured most of the sample-efficiency gain; going beyond 5 steps on the benchmarks tested in that work consistently degraded final performance. A practical stopping rule: if the synthetic-to-real update ratio exceeds 20:1 and held-out model error has not improved in the last 10 thousand environment steps, freeze the horizon or cut it in half. Log these numbers every epoch; the horizon that felt safe at 50 thousand steps may become harmful at 200 thousand if the policy visits new states the model was not trained on.
Seed model rollouts from real states, keep the horizon short, monitor held-out model error, and reduce or stop imagination when calibration deteriorates or synthetic data overwhelms the real buffer.
Synthetic transitions can quietly dominate training and pull the learner toward impossible states. If your synthetic-to-real ratio climbs without a corresponding held-out model audit, you may be optimizing on fantasy data.
Think of a navigator using dead reckoning on a foggy river: each paddle stroke introduces a small positional estimate error, and after a dozen strokes the accumulated drift can place the canoe on a sandbank that was never in the channel. The fog does not announce when the estimate stopped being useful. Imagination rollouts work the same way: each synthetic step compounds the model's small one-step error, and after enough steps the imagined state drifts into a region the real system can never occupy. The horizon is the number of strokes you trust before pulling out the GPS and checking against ground truth.
A common assumption is that longer imagination rollouts are strictly better because they generate more training data and allow the agent to "think further ahead." This is wrong in embodied AI contexts because world model error compounds with each imagined step: a model with 2 percent one-step prediction error produces an imagined state after ten steps that may lie entirely outside any region the robot can physically reach or recover from. The correct mental model is that rollout length is a trust budget, not a performance dial. Extend the horizon only as far as the model's held-out prediction error remains below a calibrated threshold, and treat any horizon beyond that threshold as dataset corruption rather than useful foresight.
For a tabletop pushing task, two or three imagined steps branched from real states may be enough to accelerate value learning. For long-horizon autonomous driving, naive long synthetic rollouts can easily invent lane states or contact events the real car would never produce.
This section connects directly to the rollout-horizon caution in Section 36.3 and to MBPO in the bibliography below.
1. Diffusion-based world models for imagination. Recent work replaces deterministic or Gaussian next-state predictors with score-based diffusion models, enabling richer multi-modal imagination over contact events. UniSim (Yang et al., 2024, Google DeepMind) trains a video diffusion model as a universal simulator of robot interactions, then uses its imagined rollouts as training data for downstream policies. The key finding is that diffusion-based imagination captures bimodal contact distributions that Gaussian ensemble models systematically blur.
2. Foundation world models and cross-embodiment imagination transfer. Large pre-trained video prediction models (e.g., Genie 2, DeepMind 2024) are being adapted as frozen priors for imagination rollouts on novel robot morphologies. The hypothesis is that a model pre-trained on internet video already encodes broad physical priors, so fine-tuning its latent space on a small amount of robot-specific data yields a world model that generalizes farther and needs fewer real trials. IRASim (Zhu et al., 2024) extends this to robot arm trajectory generation from foundation video priors.
3. Uncertainty-aware adaptive horizon scheduling. Rather than a fixed rollout length, several 2024-2025 papers propose online horizon adaptation driven by calibrated epistemic uncertainty estimates. FOWM (Feng et al., 2024) uses an ensemble disagreement signal to gate imagination steps, extending the horizon only when variance across ensemble members stays below a learned threshold. This replaces the hand-set 5 percent MSE rule with a self-calibrating schedule that adjusts to policy-induced distribution shift during training.
Open problem for a PhD student: All three directions above still treat imagination as a homogeneous stream: every state in the rollout contributes equally to policy gradient updates. A promising open problem is contact-aware selective imagination, assigning higher weight or stopping rollouts at the precise step where a contact event is predicted, because model error spikes sharply at those transitions. Combining a learned contact-event detector (or a tactile-signal predictor) with a per-step loss weighting scheme could let a robot policy extract far more useful signal from a short imagination budget, particularly for dexterous manipulation tasks where free-space prediction is cheap and contact prediction is the bottleneck.
Why is branching from replay-buffer states safer than initializing long synthetic rollouts from synthetic states created by earlier imagination? Answer using the support argument above: a real replay-buffer state is guaranteed to lie inside the region the model was trained on, while a state produced by an earlier imagined step may already have drifted outside that region, so branching a new rollout from it compounds an error that was never checked against real experience.
Imagination helps when it stays tethered to reality. Cut the tether, and the learner starts studying its own fiction.
Imagination rollouts are valuable because they multiply data use, but only when the rollout horizon is kept inside the model's trusted neighborhood.
Project Ideas
Beginner (weekend): Implement a minimal MBPO branching loop on the Gymnasium Hopper-v4 environment: train a one-step MLP dynamics model on a small replay buffer, branch short imagined rollouts at horizons 1, 3, and 5, then plot how the synthetic-to-real ratio changes the learned value function. The key challenge is setting the per-step trust threshold and observing empirically when extending the horizon stops helping and starts inflating Q-values (the learned estimates of expected future return for a given state and action).
Intermediate (1-2 weeks): Apply imagination rollouts to a contact-rich manipulation task in MuJoCo using the FetchPush-v2 environment from Gymnasium Robotics: train an ensemble dynamics model (several independently trained predictors whose spread indicates how uncertain the model is), seed imagined rollouts only from states with low ensemble disagreement, and compare final success rate against a pure model-free Soft Actor-Critic (SAC) baseline. The key challenge is that contact discontinuities make per-step model error non-uniform, so your trust metric must weight contact-phase transitions more heavily than free-space ones to avoid generating physically impossible grasps.
Design an MBPO-style training loop for a robot task. What states seed imagination, what horizon would you start with, and what metric would trigger shortening the rollout?
Step-Through: MBPO branching with concrete numbers
Trace one imagination step from a single real state. Start with replay-buffer state \(s_0 = [0.50, -0.20]\) (a 2D toy observation whose valid range is \([-1, 1]\) per dimension, so 5 percent of the range is \(0.10\)). The policy picks action \(a_0 = 0.30\). The learned model predicts a delta \(\Delta = [0.04, 0.03]\), giving imagined state \(s_1 = [0.54, -0.17]\). Held-out one-step error on this region is MSE \(= 0.004\) (per-dimension RMSE, root mean squared error, \(\approx 0.063\)), which is below the \(0.10\) threshold, so the step is trusted and kept. Now branch a second step from \(s_1\): the model predicts \(\Delta = [0.05, 0.06]\), giving \(s_2 = [0.59, -0.11]\), but here held-out RMSE \(\approx 0.12\), which exceeds \(0.10\). With \(h = 1\) you stop after \(s_1\) and add one synthetic transition \((s_0, a_0, s_1)\) to the training set. With \(h = 2\) you would also add \((s_1, a_1, s_2)\), importing a transition whose error already broke the trust budget. Same seed state, one extra step, and the dataset just absorbed a partly fictional transition.
Real-World Application: locomotion on physical robots
The MBPO branching recipe underlies TD-MPC2 deployments on quadruped and humanoid hardware, where short imagined rollouts from real replay states let the controller refine its value estimates between costly real trials. Researchers have used the same short-horizon imagination idea on the Unitree Go1 to learn locomotion gaits with a fraction of the real-world rollouts a model-free agent would need, keeping wear and tear on the legs manageable while the policy improves.
The horizon that wanted to be longer
When Janner and colleagues first scaled up MBPO, intuition said more imagined steps should mean more free data and faster learning. The experiments said the opposite: on several MuJoCo benchmarks a one-step rollout captured most of the gain, and pushing past five steps reliably made performance worse, not better. The counterintuitive lesson is that the most valuable synthetic dataset was often the shortest one. A world model is a generous liar: it will happily hand you a thousand-step trajectory, and almost none of it is true past the first few steps.
Lab: Watch the horizon turn helpful data into fiction
Goal: measure, empirically, the horizon at which imagination stops helping and starts corrupting value estimates. Tools needed: Python, gymnasium with Hopper-v4 (MuJoCo), PyTorch, and roughly 20-30 minutes including training. Setup: collect a small replay buffer (5k-10k real steps from a random or partially trained Soft Actor-Critic policy), then fit a small one-step MLP dynamics model that predicts the next observation. What to vary: the imagination horizon \(h \in \{1, 2, 3, 5, 10\}\), branching imagined rollouts from real buffer states each epoch. What to observe: (1) per-step held-out next-state MSE as a function of rollout depth, and confirm it grows roughly multiplicatively; (2) the synthetic-to-real transition ratio at each horizon; (3) the learned Q-value on a fixed set of held-out states, watching for the point where higher \(h\) inflates Q-values without improving real-environment return. Plot held-out MSE versus rollout step on one axis and final episode return versus \(h\) on another. You should see returns peak around \(h = 1\) to \(3\) and degrade past \(h = 5\), mirroring the original MBPO finding, while the per-step error curve makes the compounding mechanism visible.
Bibliography & Further Reading
Hafner, D. et al.. "Mastering Diverse Domains through World Models." (2023). https://arxiv.org/abs/2301.04104
DreamerV3 is a broad latent imagination baseline worth contrasting with explicit MBPO-style branching.
Hansen, N. et al.. "TD-MPC2: Scalable, Robust World Models for Continuous Control." (2023). https://arxiv.org/abs/2310.16828
Useful for comparing latent short-horizon planning with synthetic-data augmentation approaches.
Janner, M. et al.. "When to Trust Your Model: Model-Based Policy Optimization." (2019). https://arxiv.org/abs/1906.08253
The essential reference for short trusted imagination rollouts.