Policies Learned in Imagination › Actor-Critic Learning Inside a Model

Once we can generate rollouts, we need a learning signal. The natural choice is actor-critic, because it gives us both a critic that can reason about the tail of the horizon and an actor that can be optimized directly.

Actor-critic is itself a classical idea, and it is helpful to place it in that lineage. In approximate dynamic programming, the value function is the object of interest: policy iteration alternates between evaluating a policy and improving it greedily with respect to the value function. Actor-critic methods combine those two steps into a single loop with two function approximators, one for the value and one for the policy. Inside a learned model, the loop is unchanged in structure, but every state, reward, and transition it uses is synthetic. This is the sense in which imagination-based learning is classical control with a learned simulator swapped in: the algorithm is familiar, the data is not.

Recall from Chapter 6 of Part II that a critic $v_\xi$ estimates the expected discounted return from a state. Inside imagination, the critic is trained by regression onto multi-step return targets computed from the imaginary rewards. The standard target is the $\lambda$-return, a geometric blend of $n$-step returns that trades bias against variance. Define the temporal-difference residual along the rollout as

$$\delta_t = \hat r_t + \gamma\, v_\xi(\hat z_{t+1}) - v_\xi(\hat z_t),$$

where:

• $\delta_t$: the temporal-difference residual at step $t$, measuring the surprise of the imagined reward relative to the critic's prediction
• $\hat r_t$: the imagined reward at step $t$
• $\gamma$: the discount factor
• $v_\xi(\hat z_{t+1})$: the critic's value estimate at the next latent
• $v_\xi(\hat z_t)$: the critic's value estimate at the current latent

with the boundary condition that $\hat z_H$ is the last latent and $v_\xi(\hat z_H)$ is the bootstrap value. The $\lambda$-return is then

$$\hat G_t^{\lambda} = v_\xi(\hat z_t) + \sum_{k=t}^{H-1} (\gamma\lambda)^{k-t}\,\delta_k,$$

where:

• $\hat G_t^{\lambda}$: the $\lambda$-return target at step $t$, blending $n$-step returns to balance bias and variance
• $v_\xi(\hat z_t)$: the critic's baseline value estimate at the current latent
• $\lambda$: the blending parameter between one-step and Monte Carlo returns
• $\delta_k$: the temporal-difference residual at future step $k$
• $H$: the imagination horizon

which can be computed with a single backward pass over the rollout using the recursion $\hat G_t^{\lambda} = \delta_t + \gamma\lambda\,\hat G_{t+1}^{\lambda} + v_\xi(\hat z_t) - \gamma\lambda\, v_\xi(\hat z_{t+1})$ or, more simply, by accumulating an auxiliary variable $A_t = \delta_t + \gamma\lambda A_{t+1}$ and returning $v_\xi(\hat z_t) + A_t$. The critic loss is the mean squared error between $v_\xi(\hat z_t)$ and the detached target:

$$\mathcal{L}_{\text{critic}}(\xi) = \mathbb{E}\left[\frac{1}{H}\sum_{t=0}^{H-1}\left(\hat G_t^{\lambda} - v_\xi(\hat z_t)\right)^2\right].$$

where:

• $\mathcal{L}_{\text{critic}}(\xi)$: the critic loss as a function of the value-network parameters $\xi$
• $\mathbb{E}[\cdot]$: the expectation over imagined rollouts
• $H$: the number of rollout steps
• $\hat G_t^{\lambda}$: the detached $\lambda$-return target at step $t$
• $v_\xi(\hat z_t)$: the critic's predicted value at the latent state $\hat z_t$

In this separate critic regression, the target is detached during each update. Other algorithms deliberately differentiate through targets, so detachment is a design choice rather than a universal rule.

Why use a $\lambda$-return rather than only a one-step or full-rollout target? A one-step target relies heavily on the critic's bootstrap and can inherit its bias. A full sampled return avoids intermediate value bootstraps but may have higher sampling variance; in a deterministic model, that variance need not be large. If the rollout ends before the task does, its terminal bootstrap still carries value error. The $\lambda$-return interpolates among these targets. Its useful bias-variance balance depends on critic accuracy, model accuracy, and stochasticity rather than on a universal ordering.

For the actor there are two fundamentally different gradient estimators, and the distinction is one of the most important ideas in this chapter.

The score-function estimator. Also called REINFORCE or the likelihood-ratio estimator, this treats the sampled action as a fixed constant and reweights its log-probability by how good the outcome turned out to be:

$$\nabla_\omega \hat J_H(\omega) = \mathbb{E}\left[\sum_{t=0}^{H-1}\gamma^t\left(R_t^H-b(\hat z_t)\right)\nabla_\omega \log \pi_\omega(a_t\mid\hat z_t)\right], \qquad R_t^H=\sum_{k=t}^{H-1}\gamma^{k-t}\hat r_k.$$

where:

• $\nabla_\omega\hat J_H$: the gradient of the finite, $H$-step discounted imagined-reward objective
• $R_t^H$: sampled discounted reward-to-go from step $t$ through the truncated rollout
• $b(\hat z_t)$: a state-dependent baseline, often the critic, whose expectation against the action score is zero
• $\gamma^t$: the time weight required by the stated discounted objective
• $\log \pi_\omega(a_t \mid \hat z_t)$: the log-probability of the sampled action under the current policy

The likelihood-ratio identity is unbiased for this finite imagined objective when $R_t^H$ is sampled from the stated rollout distribution and the baseline depends on state, not the sampled action. A practical update may replace $R_t^H$ with a bootstrapped $\lambda$-return and use $\operatorname{sg}(\hat G_t^\lambda-v_\xi(\hat z_t))$ as an advantage; with an approximate critic and $\lambda<1$, that surrogate is generally biased. The estimator works with discrete actions and non-differentiable transitions. Original Dreamer used analytic gradients through latent dynamics; DreamerV3 uses a REINFORCE-style actor update.

For one sampled action, the score is weighted by how far its return lies above or below the baseline. The magnitude matters, but the estimator does not observe the derivative of reward with respect to an untried nearby action. Sampling supplies that information indirectly. A critic baseline reduces variance by subtracting predictable state value without changing the exact score identity.

The pathwise estimator. Also called the analytic or reparameterized gradient, this differentiates directly through the reward and dynamics:

$$\nabla_\omega \hat J_H(\omega) = \mathbb{E}\left[\sum_{t=0}^{H-1}\gamma^t\left(\frac{\partial \hat r_\psi}{\partial a_t}\frac{d a_t}{d\omega}+\frac{\partial \hat r_\psi}{\partial \hat z_t}\frac{d\hat z_t}{d\omega}\right)\right].$$

Here $d a_t/d\omega$ and $d\hat z_t/d\omega$ are total derivatives through all earlier sampled actions and transitions. If the actor objective includes a terminal bootstrap, its derivative must be included too. On a smooth reparameterized model, pathwise gradients often have lower variance than score-function gradients, though long Jacobian products can be unstable. They estimate the gradient of the imagined objective without requiring the model to be correct; model accuracy is needed to identify that gradient with the real-environment gradient. Discrete or discontinuous models may require a relaxation, straight-through estimator, or score-function route.

The score-function estimator uses sampled outcomes; the pathwise estimator additionally uses local model derivatives. Those derivatives can provide efficient credit assignment, but can also be misleading outside the training distribution. Neither estimator removes model bias. Their relative variance and learning speed depend on the model, horizon, policy, and optimizer.

The two can also be mixed, but their weights and bootstrap treatment define the actual objective. A REINFORCE-style bootstrapped actor surrogate, as used in later Dreamer variants, is:

$$\mathcal{L}_{\text{actor}}(\omega) = -\,\mathbb{E}\left[\sum_{t=0}^{H-1}\gamma^t\operatorname{sg}\!\left(\hat G_t^{\lambda} - v_\xi(\hat z_t)\right)\log \pi_\omega(a_t \mid \hat z_t)\right] - \eta\,\mathbb{E}\left[\sum_{t=0}^{H-1}\mathcal{H}\big[\pi_\omega(\cdot \mid \hat z_t)\big]\right].$$

The entropy coefficient $\eta$ can discourage premature concentration. Its effect depends on reward and gradient scales. The toy actor below has a learned global log-standard-deviation, so it cannot modulate exploration by state. Its optional actor.dist(s).entropy() term measures the pre-squash Normal's entropy, a surrogate regularizer rather than the entropy of bounded actions after $\tanh$. The plotted experiments set $\eta=0$ to isolate the two gradient routes.

At exactly zero variance, a point-mass policy has no ordinary log-density score; very small variance can make score estimates ill-conditioned and leave too little exploration. Entropy can prevent premature concentration, but stochasticity does not by itself protect against model exploitation.

Two possible stability tools are target critics and return normalization. A slowly updated target critic can reduce target movement; a return scale estimated independently of the current actor batch can moderate gradient magnitude across reward scales. They are not universal requirements. Our toy implementation instead uses a fixed reward scale, bounded actions, a conservative learning rate, and gradient clipping; its exact score-function branch does not normalize the sampled advantage.

Bootstrapping with function approximation can destabilize value learning, especially with off-policy data. Target networks can damp one feedback loop, while scaling changes optimizer step sizes. Neither proves convergence, and scaling the advantage using statistics of the same sampled actions can change the exact score-gradient estimator. The code below keeps that distinction visible.

Finally, note what the actor is not doing. It is not maximizing the true return. The objective is

$$\hat J(\omega) = \mathbb{E}_{\hat p_\theta,\, \pi_\omega}\left[\sum_{t=0}^{H-1}\gamma^t \hat r_\psi(\hat z_t, a_t)\right],$$

where:

• $\hat J(\omega)$: the imagined return objective as a function of policy parameters $\omega$
• $\mathbb{E}_{\hat p_\theta,\, \pi_\omega}$: the expectation taken under the learned model $\hat p_\theta$ and the current policy $\pi_\omega$
• $\gamma$: the discount factor
• $\hat r_\psi(\hat z_t, a_t)$: the model's reward prediction at latent state $\hat z_t$ and action $a_t$

an expectation under the model. Everything that follows is about the gap between $\hat J(\omega)$ and the real $J(\omega)$.

It is worth stating the mismatch in the sharpest possible terms. The world model is trained to predict observations, rewards, and transitions, which is a prediction objective. The actor is trained to maximize return under the model, which is a control objective. These two objectives are not the same, and nothing in the training procedure guarantees that the model's predictions are accurate on the states the actor will visit. This is sometimes called objective mismatch, and it is the structural reason that improving the model's average prediction error does not automatically improve the policy. A model can become better at predicting the data it was trained on while becoming worse as a training environment, if the improvement comes from fitting incidental structure that the policy then exploits.