Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

The Dreamer Family › Latent imagination and actor-critic learning

Michael BrenndoerferJuly 15, 202613 min read

World Models HandbookAll chapters

Part 8 of 12Decision-Centric Research Lineages

Chapter 45The Dreamer Family › Latent imagination and actor-critic learning

Reading section 3 of 7 · Latent imagination and actor-critic learning · 13 minSection 3 of 7 · 13 min

Section 3 of 7 · 13 min read

The Dreamer Family › Latent imagination and actor-critic learningOne section · Sections in this chapter3 / 7 ↓

Sections in this chapter7

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

We now have a model that predicts latent transitions, rewards, and continuations. Dreamer trains an actor and critic on imagined trajectories rather than unrolling the environment for each update. Depending on the generation and domain, the actor uses either analytic gradients through imagined dynamics or score-function gradients from returns while the model trajectory is stopped.

The actor update uses rewards and values predicted by the learned model rather than direct rewards from the environment. Those predictions are training targets, not ground truth. Real replay observations seed the imagined trajectories; after that, states and rewards come from the model. If the rollout leaves the data-supported region, the actor may optimize behavior that does not transfer.

The setup is as follows. At each iteration, you take a batch of real observations, encode them into latents, then roll out an imagined trajectory using the prior and the current policy for action selection. For each imagined step, we record the predicted reward, the predicted continuation probability, and the predicted value from the critic. The critic is trained against a continuation-weighted lambda-return that combines predicted rewards with critic bootstraps along the imagined trajectory and at its horizon. The actor is trained to maximize those imagined returns. This is essentially an actor-critic method, but the trajectory is entirely synthetic and lives in latent space.

Why use an actor instead of online planning? A policy amortizes decision-making across states and selects an action with one network evaluation rather than a fresh trajectory search. The actor-gradient estimator differs by generation: V1 backpropagates through imagined dynamics, reparameterizing continuous actions and using straight-through gradients for discrete actions; V2 uses REINFORCE on Atari and dynamics gradients on continuous control; V3 uses entropy-regularized REINFORCE and stops actor gradients through the world model.

An actor amortizes action selection across the states used in training. Unlike a fresh PlaNet-style CEM solve, it reuses learned parameters at the next decision. This does not guarantee smooth actions for nearby latents, and planners can also reuse computation through warm starts, learned priors, or value functions.

The toy implementation below uses a V3-style score-function update for discrete actions, but omits the cumulative imagination weights used in full Dreamer training. The world model generates states and rewards without gradients; the actor is then updated from the log-probabilities of exactly those sampled actions at exactly those decision states.

The toy actor loss maximizes expected advantage weighted by the log-probability of the chosen actions, with an entropy bonus to maintain exploration. Its unweighted surrogate is:

\mathcal{L}_\text{agent}(\theta) = -\mathbb{E}\left[\sum_{t=0}^{H-1} \left( \log \pi_\theta(a_t \mid s_t) \, \mathrm{sg}\!\left[ R^\lambda_t - V_\psi(s_t) \right] + \eta \, H[\pi_\theta(\cdot \mid s_t)] \right) \right],

where:

$\theta$ : the actor parameters
$\psi$ : the critic parameters
$\mathrm{sg}[\cdot]$ : the stop-gradient operator
$R^\lambda_t$ : the lambda-return, a weighted average of multi-step returns and bootstrapped critic predictions
$V_\psi(s_t)$ : the critic's value estimate at imagined state $s_t$
$\eta$ : the entropy coefficient
$H$ : the imagination horizon
$s_t$ : the imagined decision state before action $a_t$
$H[\pi_\theta(\cdot \mid s_t)]$ : the entropy of the actor's action distribution at $s_t$

The toy critic is trained to regress the lambda-return:

\mathcal{L}_\text{critic}(\psi) = \mathbb{E}\left[ \sum_t \left( V_\psi(s_t) - \mathrm{sg}\!\left[ R^\lambda_t \right] \right)^2 \right].

The lambda-return blends multi-step rewards with a value bootstrap at the imagination horizon. When $\lambda = 0$ it reduces to a one-step TD target. When $\lambda = 1$ , it includes rewards until a zero continuation; the horizon bootstrap is weighted by the product of intervening continuation probabilities:

R^\lambda_t = r_t + \gamma c_t \left[ (1 - \lambda) V_\psi(s_{t+1}) + \lambda R^\lambda_{t+1} \right],

where $\gamma$ is the discount factor, $c_t$ is the predicted continuation probability, and the boundary is $R^\lambda_H=V_\psi(s_H)$ . A zero continuation truncates the return; a probability between zero and one softly downweights future terms.

For an episodic Dreamer objective, imagined steps also receive cumulative weights $w_0=1$ and $w_t=\prod_{i=0}^{t-1}\gamma c_i$ . Without these weights, the actor, critic, and entropy losses would treat post-terminal imagined states as valid training data. Our toy code averages those losses across steps instead. DotWorld has no real terminal transitions and its continuation target is always one, so this example cannot test post-terminal masking. Even here, the code omits $\gamma^t$ weighting and any downweighting from imperfect continuation predictions; it is not the published objective.

Why a lambda-return rather than the simple discounted return $R_t = r_t + \gamma R_{t+1}$ ? Because the imagination horizon is finite. An unbootstrapped recursion sets the unseen boundary $R_H=0$ , so its last included target is $R_{H-1}=r_{H-1}$ and earlier near-horizon targets omit all rewards beyond the rollout. They are truncated rather than generally zero, which can bias long-horizon value estimates. Bootstrapping with the critic's own prediction at the horizon mitigates this, and $\lambda$ interpolates between low-variance (small $\lambda$ , mostly bootstrapped) and low-bias (large $\lambda$ , mostly Monte Carlo) returns. This is the same TD( $\lambda$ ) idea that underlies much of modern deep RL.

Smaller $\lambda$ relies more heavily on the critic; larger $\lambda$ relies more heavily on accumulated model rewards. The familiar bias-variance description is conditional because model error and critic error both contribute. A value of 0.95 is an empirical setting, not a universal optimum.

The actor objective depends on both its action distribution and the action-conditioned latent trajectory. When every stochastic node on that path is reparameterizable, differentiating through the model gives a pathwise gradient of the learned-model objective; it can still be biased relative to real-environment return when the model is inaccurate and is often lower variance than REINFORCE. DreamerV2's categorical straight-through path is instead a biased surrogate estimator, and its Atari actor uses REINFORCE. DreamerV3 stops world-model gradients and uses a REINFORCE-style signal weighted by the critic's advantage.

Stopping the pathwise derivative does not prevent exploitation. A score-function update can still reinforce actions whose modeled transitions or rewards produce spuriously high returns. REINFORCE removes one gradient path and was reported to work across DreamerV3's domains; it does not make either learned head trustworthy outside the data distribution.

Now let's implement the actor and critic on top of our trained RSSM. We add a small MLP actor that takes the latent state and outputs a categorical action distribution, plus a critic that takes the latent state and outputs a scalar value.

In[16]:

Code

class Actor(nn.Module):
    def __init__(
        self, latent_dim=16, deter_dim=64, action_dim=4, hidden_dim=128
    ):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(deter_dim + latent_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, action_dim),
        )

    def forward(self, deter, latent):
        x = torch.cat([deter, latent], dim=-1)
        return torch.distributions.Categorical(logits=self.net(x))


class Critic(nn.Module):
    def __init__(self, latent_dim=16, deter_dim=64, hidden_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(deter_dim + latent_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, deter, latent):
        x = torch.cat([deter, latent], dim=-1)
        return self.net(x).squeeze(-1)

The rollout records each pre-action decision state, sampled action, predicted reward, and successor state. The world model is frozen, matching the stopped-dynamics estimator used by V3 and by this toy discrete-action example.

In[17]:

Code

@torch.no_grad()
def encode_seed(model, obs_seed):
    B = obs_seed.shape[0]
    device = obs_seed.device
    deter = model.initial_state(B, device)
    prev_action = torch.zeros(B, model.action_dim, device=device)
    prev_latent = torch.zeros(B, model.latent_dim, device=device)
    deter, (_, _), (q_mean, _) = model.forward_step(
        prev_action, prev_latent, deter, obs_seed.view(B, -1)
    )
    return deter, q_mean


@torch.no_grad()
def imagine_trajectory(model, actor, obs_seed, horizon=15, gamma=0.99):
    deter, latent = encode_seed(model, obs_seed)
    device = obs_seed.device
    B = obs_seed.shape[0]

    latents, deters, actions = [], [], []
    rewards, continues = [], []

    for _ in range(horizon):
        deters.append(deter)
        latents.append(latent)
        dist = actor(deter, latent)
        action = dist.sample()

        action_onehot = F.one_hot(action, num_classes=model.action_dim).float()
        deter, (p_mean, _), _ = model.forward_step(
            prev_action=action_onehot, prev_latent=latent, deter=deter, obs=None
        )
        latent = p_mean
        _, rew_hat, cont_hat = model.decode(deter, latent)

        actions.append(action)
        rewards.append(rew_hat)
        continues.append(cont_hat)

    return dict(
        deters=torch.stack(deters, dim=1),
        latents=torch.stack(latents, dim=1),
        actions=torch.stack(actions, dim=1),
        rewards=torch.stack(rewards, dim=1),
        continues=torch.stack(continues, dim=1),
        final_deter=deter,
        final_latent=latent,
        gamma=gamma,
    )

The imagine_trajectory function is the core of Dreamer. Notice that the model is asked for its prior at every step, never its posterior, because there is no new observation to condition on. The action is sampled from the current actor, the reward and continuation are read from the model's heads, and the latent is carried forward. The whole imagined episode is generated by the model playing both the world and the reward function, with the actor supplying the choices.

Now we compute the lambda-returns from the imagined rewards, continuation predictions, and critic values at each step. We use backward recursion, which is the cleanest way to compute and it avoids the recursion depth growing with the horizon.

In[18]:

Code

def compute_lambda_returns(rewards, continues, values, gamma=0.99, lam=0.95):
    # rewards and continues: (B, H); values: (B, H + 1)
    B, H = rewards.shape
    returns = torch.zeros_like(rewards)
    next_return = values[:, H]
    for t in reversed(range(H)):
        next_value = values[:, t + 1]
        bootstrap = (1.0 - lam) * next_value + lam * next_return
        next_return = rewards[:, t] + gamma * continues[:, t] * bootstrap
        returns[:, t] = next_return
    return returns


def actor_critic_loss(
    model,
    actor,
    critic,
    obs_seed,
    horizon=15,
    gamma=0.99,
    lam=0.95,
    entropy_coef=1e-3,
):
    traj = imagine_trajectory(
        model, actor, obs_seed, horizon=horizon, gamma=gamma
    )

    deters = traj["deters"]
    latents = traj["latents"]
    rewards = traj["rewards"]
    continues = traj["continues"]

    B, H, D = deters.shape
    flat_deter = deters.reshape(B * H, D)
    flat_latent = latents.reshape(B * H, -1)
    decision_values = critic(flat_deter, flat_latent).reshape(B, H)
    final_value = critic(traj["final_deter"], traj["final_latent"]).unsqueeze(1)
    values = torch.cat([decision_values, final_value], dim=1)

    returns = compute_lambda_returns(
        rewards, continues, values.detach(), gamma=gamma, lam=lam
    )
    advantage = (returns - decision_values).detach()

    # log_probs must be recomputed FROM the sampled actions so gradients flow to the actor;
    # the current log_probs come from @torch.no_grad() imagine_trajectory and carry no grad_fn.
    dist_for_loss = actor(flat_deter, flat_latent)
    log_probs = dist_for_loss.log_prob(traj["actions"].reshape(B * H)).reshape(
        B, H
    )
    actor_loss = -(log_probs * advantage).mean()

    critic_loss = F.mse_loss(decision_values, returns.detach())

    dist = actor(flat_deter, flat_latent)
    entropy = dist.entropy().mean()
    total = actor_loss + critic_loss - entropy_coef * entropy
    return total, dict(
        actor=actor_loss.item(),
        critic=critic_loss.item(),
        entropy=entropy.item(),
        return_mean=returns.mean().item(),
    )

The actor_loss line implements the toy's unweighted REINFORCE-with-advantage rule: the log-probability of the sampled action is multiplied by the advantage and used as a loss to maximize. The advantage is detached, so the actor is not tempted to move the value estimate itself. It can only change the action distribution. The critic_loss is a plain mean-squared error against the detached lambda-return. The entropy term is subtracted from the total loss because it is added to the objective, encouraging the actor to keep its action distribution spread out early in training. For an episodic task, the actor, critic, and entropy reductions would need the cumulative discount-and-continuation weights described above.

Now we train the actor and critic on imagined trajectories. Note that the model is frozen during this loop: only the actor and critic update. In a full Dreamer implementation, you would alternate between model updates and agent updates. You would also collect new real data between iterations. Here we keep the model fixed so we can isolate the actor-critic mechanism.

In[19]:

Code

actor = Actor()
critic = Critic()
for parameter in model.parameters():
    parameter.requires_grad_(False)
agent_optimizer = torch.optim.Adam(
    list(actor.parameters()) + list(critic.parameters()), lr=3e-4
)

seed_pool = torch.tensor(next_obs[:256], dtype=torch.float32)

agent_history = []
for step in range(400):
    idx = np.random.randint(0, len(seed_pool), size=16)
    seeds = seed_pool[idx]
    loss, parts = actor_critic_loss(model, actor, critic, seeds, horizon=15)
    agent_optimizer.zero_grad()
    loss.backward()
    nn.utils.clip_grad_norm_(
        list(actor.parameters()) + list(critic.parameters()), 100.0
    )
    agent_optimizer.step()
    if step % 50 == 0:
        agent_history.append(
            (step, parts["actor"], parts["critic"], parts["return_mean"])
        )

Out[20]:

Console

Imagined return: -0.9208 -> -15.8317

Several train/evaluation differences matter here. seed_pool takes only the first 256 of 2,600 training transitions as isolated replay frames, so encode_seed starts RSSM memory at zero for each imagined trajectory; real-environment evaluation starts fresh episodes and carries recurrent memory through the steps of each episode. The actor samples actions during imagination, while evaluation below selects the most probable action. These seed-distribution, memory, and action-selection differences can affect real performance, so comparing the two return measurements does not isolate world-model exploitation.

The mean imagined lambda-return is a changing model-side diagnostic, not real environment return. Because its critic bootstrap changes during training, movement in this quantity cannot be attributed to the actor alone. In this run it decreases while the trained policy also performs worse than the random baseline in the real environment. These results do not establish an imagination-to-reality transfer failure or identify why the policy underperforms. Model exploitation, estimator noise, optimization failure, inadequate state learning, and the seed-state and action-selection mismatches just described remain possibilities. In a full system, poor real performance motivates collecting fresh data and updating the model.

Let's evaluate the trained actor greedily on the real environment, unrolling it step by step. Does the actor, trained on imagined rollouts seeded from real observations, keep the dot near the center in the real world?

In[21]:

Code

@torch.no_grad()
def evaluate_actor(model, actor, world, episodes=20, steps=50, image_size=16):
    rng = np.random.default_rng(999)
    total_rewards = []
    for _ in range(episodes):
        pos, vel = world.reset(rng)
        episode_reward = 0.0
        deter = model.initial_state(1, torch.device("cpu"))
        prev_action = torch.zeros(1, model.action_dim)
        prev_latent = torch.zeros(1, model.latent_dim)
        for _ in range(steps):
            img = world.render(pos, image_size)
            obs_t = torch.tensor(img, dtype=torch.float32).unsqueeze(0)
            deter, _, (q_mean, _) = model.forward_step(
                prev_action, prev_latent, deter, obs_t.view(1, -1)
            )
            latent = q_mean
            dist = actor(deter, latent)
            action = int(dist.probs.argmax(dim=-1).item())
            pos, vel, reward, done = world.step(pos, vel, action, rng)
            episode_reward += reward
            prev_action = one_hot_action(
                torch.tensor([action]), model.action_dim, torch.device("cpu")
            )
            prev_latent = latent
        total_rewards.append(episode_reward)
    return np.mean(total_rewards), np.std(total_rewards)


## Compute the random-policy baseline for comparison
rng = np.random.default_rng(777)
random_episode_rewards = []
for _ in range(20):
    pos, vel = world.reset(rng)
    r = 0.0
    for _ in range(50):
        action = rng.integers(0, 4)
        pos, vel, reward, done = world.step(pos, vel, action, rng)
        r += reward
    random_episode_rewards.append(r)
random_mean = np.mean(random_episode_rewards)
random_std = np.std(random_episode_rewards)

trained_mean, trained_std = evaluate_actor(model, actor, world)
print(f"Random return: {random_mean:.4f} +/- {random_std:.4f}")
print(f"Trained return: {trained_mean:.4f} +/- {trained_std:.4f}")

This comparison is diagnostic rather than guaranteed. Improvement would support transfer on this toy run; a worse return can reflect model exploitation, an inadequate model, estimator noise, optimization failure, or the seed-state and action-selection differences between training and evaluation. In practice, Dreamer alternates data collection, model learning, and actor-critic learning. The benchmark systems also include architectural, optimization, replay, and interface details beyond scaling networks and training time.

The most important takeaways from this section:

Dreamer trains the actor and critic in imagination, using the learned model as a latent simulator.
The actor's signal comes from the critic's advantage, not from real rewards at training time.
The critic is trained with a lambda-return whose continuation factor truncates at zero and otherwise downweights future terms.
This toy implementation and DreamerV3 stop actor gradients through the world model; V1 and V2 use the generation-specific estimators described above.
The imagined return and the real return can diverge. Closing that gap is what the rest of the Dreamer pipeline is for.