Policies Learned in Imagination › Latent Imagined Rollouts

The first design question is what the rollout operates on. A generative pixel-level model can roll forward in observation space if it predicts action-conditioned future images and rewards. But decoding images can be expensive, while the critic may need only a compact, task-relevant summary. Reconstruction can still be useful supervision for learning that summary even when decoding is unnecessary during behavior learning.

This is where notions of model quality separate. Predictive fidelity asks whether the next-step prediction matches reality. Rollout fidelity asks whether predictions remain useful when fed back over many steps. Representation quality asks whether the latent retains information needed for control. Uncertainty quality asks whether the model signals its limitations. Decision usefulness asks whether a policy trained against it works in the real environment. These properties can diverge. Latent imagination avoids repeated decoding during behavior learning; it does not make the other properties optional.

Latent imagination sidesteps this. Building on the recurrent and transformer state-space architectures from Part V: World-Model Architectures, the world model provides an encoder, a latent transition, and a reward head. The rollout happens entirely inside the latent space:

where:

• $o_{\le t}$: the observation history up to and including time $t$, which the encoder consumes to produce a latent
• $a_{<t}$: the action history before time $t$, included so the encoder can account for how the current observation was reached
• $q_\phi$: the encoder, parameterized by $\phi$, that maps observations and past actions into the latent state $\hat z_t$
• $p_\theta$: the latent transition model, parameterized by $\theta$, which defines a distribution over the next latent state given the current latent and action
• $r_\psi$: the reward head, parameterized by $\psi$, that predicts the scalar reward from a latent state and action
• $\hat z_t$: the inferred latent state at time $t$, intended to summarize the task-relevant past; whether it succeeds is empirical
• $\hat r_t$: the predicted reward at time $t$, used as the synthetic training signal for the actor and critic

The equations resemble a classical control pipeline. The encoder $q_\phi$ is a learned state estimator, though it need not produce a calibrated sufficient belief. The transition model $p_\theta$ predicts how that state evolves under control. The reward head $r_\psi$ predicts a task reward supplied by the environment; it is not itself a learned preference specification. The actor and critic play the roles of controller and value estimate. Under partial observability, history or a belief estimate is generally useful, although a reactive observation policy can suffice in special tasks.

No observation needs to be decoded during these imagined actor and critic updates. The actor and critic read $\hat z_t$, while the reward head scores $(\hat z_t,a_t)$. A decoder may still be important when learning the representation.

There are three practical reasons this matters. First, skipping a high-dimensional decoder can reduce rollout compute, depending on architecture and observation resolution. Second, gradients need not pass through a pixel decoder to reach the policy. Third, reconstruction can pressure a model to preserve visual detail irrelevant to action selection; reward and transition objectives can additionally shape the latent toward control.

The representation question is which information the controller needs. A latent can reconstruct fine visual detail yet omit a small cue that changes the right action; another can reconstruct imperfectly but retain that cue. Neither ordering is guaranteed by the loss alone. Reward and value heads can help shape what the actor extracts from the latent, while the model's training objective still determines what information reaches those heads.

The rollout is then a single vectorized loop. Starting from a batch of $B$ latent states, at each of $H$ steps:

• Sample an action $a_t \sim \pi_\omega(\cdot \mid \hat z_t)$, where $\pi_\omega$ is the policy parameterized by $\omega$.
• Step the transition model to get $\hat z_{t+1}$.
• Score the reward $\hat r_t$ and query the critic for $\hat v_t = v_\xi(\hat z_t)$, where $v_\xi$ is the value function parameterized by $\xi$.

Because the $B$ rollouts are independent, each step can be batched on a GPU. There are still $H$ sequential transition steps and arithmetic scaling with $BH$; batching reduces wall-clock overhead, not the amount of work. This is the concrete form of the amortization bargain: training is expensive in aggregate but parallelizable across starts.

After the loop, query the critic once more at the final latent to get a bootstrap value $\hat v_H$. With $\gamma=0.99$ and $H=20$, that terminal value receives weight $\gamma^{20}\approx0.818$. For an infinite stream of constant rewards, about 82% of the normalized discounted mass lies after the first 20 rewards. In a continuing task, omitting the bootstrap can therefore undervalue delayed payoffs.

The scale $1/(1-\gamma)=100$ describes the geometric discount tail, not a literal cutoff. The first 20 rewards contain $1-0.99^{20}\approx18.2\%$ of the normalized mass for a constant stream. A terminal value estimate, as in finite-horizon dynamic programming or receding-horizon control, represents the continuing value beyond the sampled prefix. Its usefulness depends on its accuracy.

Whether the loop is differentiable matters. If the latent transition is a continuous distribution, we can sample with the reparameterization trick, $\hat z_{t+1} = \mu_\theta(\hat z_t, a_t) + \sigma_\theta(\hat z_t, a_t) \odot \epsilon$ with $\epsilon \sim \mathcal{N}(0, I)$, where $\mu_\theta$ and $\sigma_\theta$ are the mean and standard deviation predicted by the latent transition model and $\odot$ denotes elementwise multiplication, and gradients flow straight through the model. If the latent is discrete, as in many categorical state-space models, the sample is not differentiable. The usual fixes are a straight-through estimator, which passes the gradient of a relaxed sample through a hard sample, or a gradient-free actor update that treats the dynamics as a black box. Each choice changes what the actor gradient estimator can do, which we take up next.

Differentiability determines whether an exact pathwise derivative through sampled transitions is available. Continuous reparameterized models support it directly. Discrete or discontinuous models may use a relaxation, straight-through approximation, smoothing, or score-function estimator; none is automatically robust to model error. The choice affects both gradient bias or variance and the model derivatives the actor can exploit.

If actor gradients could update an unconstrained reward or transition head, the optimizer could improve the reported outcome without improving behavior. That shortcut is a risk, not an inevitable unbounded trajectory: predictive losses, regularization, and head ranges can counter it. In this chapter's code the model is fitted first and frozen during actor learning, so the two objectives remain separate.