Deep Dive into Dreamer V1 Architecture
Dreamer was introduced in 2019 by Danijar Hafner and team at Google Brain and DeepMind as a landmark model-based reinforcement learning (MBRL) architecture
Dreamer was one of the first RL model that showed strong generalizability across diverse RL tasks and environments, with a single configuration achieving competitive performance across more than 150 tasks without task-specific tuning. It also demonstrated the ability to solve challenging long-horizon, sparse-reward problems, including collecting diamonds in Minecraft from scratch without human data and supervision
In this article, we’ll deep dive into the core architecture of the first Dreamer V1 (2019), breaking down the architecture and math behind Recurrent State-Space Model (RSSM) and latent actor-critic framework, and walk step-by-step through our PyTorch code implementation.
Dreamer Architecture: High-Level Overview
At a high level, the Dreamer architecture does two things: learns latent environment dynamics and behavior that maximizes reward.
The latent environment is an unknown distribution that emits observations and rewards in response to agent actions. In a video game, for example, the actions are the player’s controls, and the observations are the new frames the game renders in response.
Learning the latent environment’s dynamics amounts to understanding how the environment responds to the agent’s actions: what new states it produces. One way to do this is to predict a future state from the current observation, the previous state, and the action taken. Let’s take a video game example. We could use a set of game frames and actions and try to learn a temporal deep neural model to predict next states generating new game frames.
This is exactly what Dreamer’s representation model does: it infers the current latent state from the previous state, the action taken, and the new observation. But Dreamer goes further. It can predict how the environment would behave without interacting with it: acting purely in imagination. This mirrors the way people mentally rehearse the outcomes of their actions before taking them, or how chess players play out move sequences and variations entirely in their heads.
In Dreamer, this is handled by the transition model, which predicts the next latent state from the previous state and action alone, with no observation involved. That’s precisely what makes pure imagination possible: latent imagination uses the transition model to unroll imagined trajectories forward from real states, entirely within the learned latent space, with no need to touch the environment.
The second component of Dreamer is behavior learning. Dreamer v1 takes an actor-critic approach, learning behavior entirely within the imagined world model. The actor learns which actions maximize returns, while the critic estimates the value of each imagined state; together, they form the behavior-learning machinery that operates purely in imagination.
The dynamics and behavior components feed into each other in a continuous loop: the world model is trained on real experience, imagination is generated within it, and the actor-critic is trained on that imagined experience. With this overall picture in mind, the next section walks through each stage in detail, starting with dynamics learning.
World model and latent dynamics
Latent dynamics is learned by a group of models collectively referred to as the “world model.” These are the representation, transition, reward, and observation models. Let’s look at how they work together.
Representation model
The representation model’s (also called the posterior model) job is to predict the next environment state given the previous latent state $s_{t-1}$, previous action $a_{t-1}$, and current observation $o_t$.
\[p(s_t \mid s_{t-1}, a_{t-1}, o_t).\]At each step, the representation model outputs a posterior Gaussian distribution, from which the posterior state is sampled.
The representation model is implemented as a Recurrent State-Space Model (RSSM) model with a CNN observation encoder.
The RSSM (Recurrent State Space Model)
The deterministic component acts as the “memory” of the representation model tracking what has happened across a sequence and preserving long-term trends and temporal dependencies. Crucially, as you see from the formula, the deterministic state does not depend on the current observation directly, which is exactly what allows the deterministic path to simulate forward trajectories without need of an observation at all.
The stochastic component $z_t$ captures uncertainty of the state outcome and hidden latent factors of the environment. It’s represented as a latent variable sampled from a Gaussian distribution conditioned on the deterministic state $h_t$ and encoded obsevation $e_t$. This means that rather than maintaining a single hidden state, the model keeps a probability distribution over possible world states at each timestep, and samples from it.
\[q_\phi(s_t \mid h_t, e_t) = \mathcal{N}\left(s_t \mid \mu_\phi(h_t, e_t), \, \text{diag}(\sigma_\phi^2(h_t, e_t))\right)\]Here $h_t$ is the deterministic state, $e_t$ is the encoded observations and $\mu_\phi(h_t, e_t)$ and $\sigma_\phi^2(h_t, e_t)$ are the outputs of a neural network that computes the mean and variance parameters of the Gaussian distribution.
In short, the deterministic path provides temporal continuity and memory; the stochastic path captures uncertainty about the true environment state: things the agent cannot know for certain from its observations alone.
Latent dynamics learning is implemented in the train_world_model method, which encodes the observations and passes them to the RSSM’s observe method:
# excerpt from Dreamer.train_world_model
obs = self.preprocess(batch.observation)
with self._autocast():
embed = self.encoder(obs)
post, prior = self.rssm.observe(embed, batch.action, batch.is_first)
# RSSM.observe
def observe(
self, embed: torch.Tensor, action: torch.Tensor, is_first: torch.Tensor
) -> tuple[State, State]:
"""Filter a batch of sequences (B, T, ...). ``action[:, t]`` is the action that led to step t; the state is reset wherever ``is_first`` is set."""
batch, length = embed.shape[:2]
state = self.initial(batch, embed.device)
posts, priors = [], []
for t in range(length):
keep = (1.0 - is_first[:, t]).unsqueeze(-1)
state = state.map(lambda x: x * keep)
state, prior = self.obs_step(state, action[:, t] * keep, embed[:, t])
posts.append(state)
priors.append(prior)
return State.stack(posts, 1), State.stack(priors, 1)
In the code above, we begin with preprocessing our batch of observations and encoding it with the CNN encoder that outputs a flat high-dimensional vector. Inside observe, we also initialize the state of the first step (both its stochastic and deterministic parts) with all zeros.
We then generate a latent dynamics sequence using representation model: at each step, obs_step advances the state using the action that led to that step and the encoded observation, and the state is reset to zeros wherever a new episode starts (is_first). obs_step begins by calling img_step:
def img_step(self, prev: State, action: torch.Tensor) -> State:
"""Prior step: advance the state with an action, without an observation."""
x = self.img_in(torch.cat([prev.stoch, action], -1))
deter = self.gru(x, prev.deter).float()
mean, std = self._stats(self.prior_net(deter))
stoch = mean + std * torch.randn_like(std)
return State(mean, std, stoch, deter)
First, we compute the deterministic state of the RSSM by passing the previous deterministic state, together with the concatenated action and previous stochastic state, through a Gated Recurrent Unit (GRU). In the same call, img_step passes the new deterministic state through an MLP to predict the prior over the stochastic state; this is the transition model, which we return to below.
Once we have $h_t$, we can compute the new stochastic state $z_t$. This is handled by the RSSM’s obs_step method, which takes the deterministic state computed above and concatenates it with the encoded observation.
def obs_step(
self, prev: State, action: torch.Tensor, embed: torch.Tensor
) -> tuple[State, State]:
"""Posterior step: returns (posterior, prior) for the next time step."""
prior = self.img_step(prev, action)
mean, std = self._stats(self.post_net(torch.cat([prior.deter, embed], -1)))
stoch = mean + std * torch.randn_like(std)
return State(mean, std, stoch, prior.deter), prior
The result is passed through an MLP that predicts the parameters of a normal distribution, from which we sample the posterior using the reparameterization trick. This posterior then becomes the previous state for the next step of the loop. obs_step returns the prior alongside the posterior, because the KL term of the loss needs both.
Note
Here, the term posterior is used because we have a distribution after having observed the data.
Transition model
The transition model is similar to the representation model in that it also predicts the next state of the environment, but it does so without looking at observations. Instead, it relies on the recurrent deterministic state (which already incorporates the previous action and stochastic state) to forecast the prior distribution: $q(s_t \mid s_{t-1}, a_{t-1})$.
Intuitively, the transition model makes it possible to generate trajectories purely in latent space, with no feedback from real-world observations. This capability is essential for latent imagination during behavior learning. Its robustness comes from minimizing the KL divergence against the representation model (the posterior), which is grounded in the observation signal.
Structurally, the transition model is an integral part of the RSSM architecture: it is the img_step method shown above. The deterministic hidden state is passed through an MLP to predict the parameters of a normal distribution, from which the stochastic prior state is sampled using the reparameterization trick. As we’ll see, this transition function plays a key role in generating multi-step imagined rollouts based entirely on previous hidden states and policy actions.
The world model has additional components needed to complete latent dynamics training: the reward and observation models, and an optional continue model.
Reward Model
The reward model predicts the scalar reward from the current latent state, $q(r_t \mid s_t)$, providing the signal the agent ultimately optimizes. It is implemented as an MLP that parameterizes a normal distribution over expected rewards, predicting the scalar mean $\hat{r}_t$ while holding the standard deviation fixed at $\sigma = 1.0$.
class NormalHead(nn.Module):
"""Unit-variance Gaussian over a scalar; used for reward and value."""
def __init__(self, feat_dim: int, units: int, layers: int, activation: str):
super().__init__()
self.mlp = MLP(feat_dim, 1, units, layers, activation)
def forward(self, feat: torch.Tensor) -> D.Normal:
return D.Normal(self.mlp(feat).squeeze(-1).float(), 1.0)
Observation Model
The observation model reconstructs the original observation from the latent state, $q_\theta(o_t \mid s_t)$, and serves as a training signal. By checking how well the latent state can recover the real observation, the model is pushed to encode the environment correctly.
It’s implemented as a decoder that reconstructs high-dimensional observations (e.g video game frames) from the world model’s latent state $[h_t, s_t]$. Rather than directly predicting fixed pixel values like a deterministic autoencoder, it parameterizes a conditional probability distribution over the image space. In Dreamer V1 this is a diagonal Gaussian with mean $\mu_t$ predicted by the network and fixed identity variance.
In the code implementation, the flat concatenated latent vector is first projected via a dense layer into a 3D tensor, then progressively upsampled to the original image dimensions using transposed convolutions.
class Decoder(nn.Module):
"""Transposed-conv decoder: latent features -> Normal(image, 1) over (C, H, W)."""
def __init__(self, feat_dim: int, obs_shape: tuple[int, int, int], config: ModelConfig):
super().__init__()
act = getattr(nn, config.cnn_act)
depth = config.cnn_depth
self.obs_shape = obs_shape
self.fc = nn.Linear(feat_dim, 32 * depth)
channels = [32 * depth, 4 * depth, 2 * depth, depth, obs_shape[0]]
layers: list[nn.Module] = []
for i, kernel in enumerate(config.decoder_kernels):
layers.append(nn.ConvTranspose2d(channels[i], channels[i + 1], kernel, stride=2))
if i < len(config.decoder_kernels) - 1:
layers.append(act())
self.deconv = nn.Sequential(*layers)
with torch.no_grad():
out = self.deconv(torch.zeros(1, 32 * depth, 1, 1)).shape[1:]
if tuple(out) != tuple(obs_shape):
raise ValueError(
f"Decoder produces {tuple(out)} but observations are {obs_shape}; "
"adjust model.decoder_kernels or env.size"
)
def forward(self, feat: torch.Tensor) -> D.Distribution:
batch_shape = feat.shape[:-1]
x = self.fc(feat).reshape(-1, self.fc.out_features, 1, 1)
mean = self.deconv(x).reshape(*batch_shape, *self.obs_shape).float()
return D.Independent(D.Normal(mean, 1.0), len(self.obs_shape))
Continue Model
Another component is the continue model (or discount predictor). Instead of relying on a fixed global discount factor, it predicts whether an environment sequence will continue at a given step ($\gamma_t = 1$) or terminate ($\gamma_t = 0$). The continue model $p_\psi(\gamma_t \mid s_t)$ is implemented as an MLP that parameterizes a Bernoulli distribution over the continuation probability $\hat{\gamma}_t \in [0, 1]$ given the latent state $s_t$. In our implementation it is optional: it is switched off for DeepMind Control, where episodes only end by time limit and the fixed discount is used, and switched on for Atari.
With this, we have all the ingredients needed to optimize the world model.
Latent dynamics loss function
Dreamer V1 uses the information bottleneck objective proposed by Tishby et al.
where $\beta$ is a scalar and $i_t$ are dataset indices that determine the observations
\[p(o_t \mid i_t) = \delta(o_t - \bar{o}_t)\]Intuitively, the first term $\max I(s_{1:T}; (o_{1:T}, r_{1:T}) \mid a_{1:T})$ encourages the latent states to retain information useful for predicting observations and rewards. The latent states must be good enough to reconstruct what’s been seen and predict what reward follows. This forces the model to actually encode useful information about the environment.
The second term dictates the opposite: don’t greedily extract everything from the current observation. It penalizes the model for relying too heavily on what it sees right now.
The tension between them forces the model to develop memory. If you can’t just use the current frame as a ground truth, you have to carry information forward through time from previous steps. The result is a model that learns strong long-term dependencies. The scalar β controls the trade-off: higher β means stronger pressure to compress, forcing the model to rely more on its recurrent memory and less on each new observation.
By expanding the components and maximizing the evidence lower bound (ELBO), or more generally, the Variational Information Bottleneck (VIB), we obtain an objective function that jointly optimizes the components of the world model. Specifically, it combines log-likelihood reconstruction terms for the observation model $q(o_t \mid s_t)$ and reward model $q(r_t \mid s_t)$ with a KL-divergence regularizer that penalizes the mismatch between the posterior representation model $p(s_t \mid s_{t-1}, a_{t-1}, o_t)$ and the prior transition dynamics $q(s_t \mid s_{t-1}, a_{t-1})$:
\[J_{REC} = \mathbb{E}_{\sim p, \theta} \left[ \sum_t \left( J_{O_t} + J_{R_t} + J_{D_t} \right) \right] + \text{const}\] \[J_{O_t} = \ln q(o_t \mid s_t)\] \[J_{R_t} = \ln q(r_t \mid s_t)\] \[J_{D_t} = -\beta \, \text{KL} \left( p(s_t \mid s_{t-1}, a_{t-1}, o_t) \,\parallel\, q(s_t \mid s_{t-1}, a_{t-1}) \right)\]Beyond modeling what the environment looks like and the rewards it yields, the agent also needs to know when the episode ends. To achieve this, we include an auxiliary continue model loss, which is trained using Binary Cross-Entropy (BCE) loss against the ground-truth status $\gamma_t = 1 - \text{terminal}_t$:
\[\mathcal{L}_{\text{continue}}(\psi) = - \left[ \gamma_t \log \hat{\gamma}_t + (1 - \gamma_t) \log (1 - \hat{\gamma}_t) \right]\]While the information bottleneck and reconstruction terms ensure the latent space captures rich spatial and reward dynamics, the continue model ensures the latent states also encode temporal boundaries. High loss penalizes misclassifying terminal states (such as predicting continuation when an episode ends).
The derivation of this objective is in the Additional Material section at the end.
Now that we have our loss function, let’s compute its components. For the reconstruction term (first term of the information bottleneck) we need the outputs of the reward and observation model.
The observation model or decoder is trained by maximizing the log-likelihood of the observed images under the predicted distribution (which, under a unit-variance Gaussian assumption, simplifies to minimizing the Mean Squared Error between reconstructed and target images):
image_loss = -self.decoder(feat).log_prob(obs).mean()
The reward loss is computed via negative log-likelihood (NLL):
\[\mathcal{L}_{\text{reward}}(\phi) = -\log q_\phi(r_t \mid s_t) \propto \frac{1}{2} (r_t - \hat{r}_t)^2\]Because the standard deviation is fixed to 1.0, this NLL loss is directly proportional to standard Mean Squared Error (MSE).
The continue model loss, is added as the negative Bernoulli log-likelihood (i.e. BCE) of the ground-truth status $\gamma_t = 1 - \text{terminal}_t$, weighted by continue_scale.
if self.cont is not None:
cont_loss = -self.cont(feat).log_prob(1.0 -batch.is_terminal).mean()
loss = loss + t.continue_scale * cont_loss
Finally, there’s the bottleneck term $J_{D_t}$, implemented via the KL-divergence regularizer:
\[\mathcal{L}_{\text{KL}}(\theta) = \beta \, \mathbb{E}_{p} \left[ \sum_t \text{KL} \left( p(s_t \mid s_{t-1}, a_{t-1}, o_t) \,\parallel\, q(s_t \mid s_{t-1}, a_{t-1}) \right) \right]\]Full world model training code:
def train_world_model(self, batch: Batch) -> tuple[State, dict[str, torch.Tensor]]:
t = self.config.train
obs = self.preprocess(batch.observation)
with self._autocast():
embed = self.encoder(obs)
post, prior = self.rssm.observe(embed, batch.action, batch.is_first)
feat = post.feat
image_loss = -self.decoder(feat).log_prob(obs).mean()
reward_loss = -self.reward(feat).log_prob(self.transform_reward(batch.reward)).mean()
kl_loss, kl = self.rssm.kl_loss(post, prior, t.free_nats)
loss = t.kl_scale * kl_loss + image_loss + reward_loss
metrics = {}
if self.cont is not None:
cont_loss = -self.cont(feat).log_prob(1.0 - batch.is_terminal).mean()
loss = loss + t.continue_scale * cont_loss
metrics["continue_loss"] = cont_loss.detach()
grad_norm = self._optimize(
"model", self.model_opt, loss, [p for mod in self.world_model for p in mod.parameters()]
)
...
Behavior Training
The task of the behavior model is to learn an optimal policy for acting in the RL environment. To achieve this, Dreamer V1 uses an actor-critic framework consisting of an action model (actor) and a value model (critic).
What distinguishes Dreamer V1 from traditional approaches is that the policy is learned entirely within the imagined trajectories of the world model. The direct benefit is that the agent does not need to interact as frequently with the real environment to generate trajectories. Furthermore, this enables the application of Dreamer to domains like robotics, where real-world training carries severe sample-efficiency and safety limitations.
The high-level flow proceeds as follows:
- Imagination Rollout: We generate imagined trajectories starting from world states with horizon $H$ by sampling actions from the action model and predicting subsequent latent states via the transition model.
- Value Estimation: The value model predicts the value of each imagined state. These predictions are regressed toward state-value estimates computed via $\lambda$-return (using predicted values as bootstraps and immediate rewards to ground the estimates).
- Cooperative Training: The action and value models are trained cooperatively in a policy-iteration fashion. The action model maximizes value estimates, while the value model aligns its predictions with the computed state-value targets.
Actor (Action Model)
The task of the action model is to maximize the estimated value. The action distribution depends on the imaginary state:
\[a_\tau \sim q_\phi(a_\tau \mid s_\tau)\]The actor is implemented as a dense network that takes the current imaginary state and outputs an action distribution. For continuous actions, this distribution is a tanh-transformed Gaussian to ensure actions remain bounded within $[-1, 1]$; for discrete actions, it is a one-hot categorical. To enable backpropagation through stochastic sampling, the action itself is sampled using the reparameterization trick:
\[a_\tau = \tanh\!\big(\mu_\phi(s_\tau) + \sigma_\phi(s_\tau)\,\varepsilon\big), \quad \varepsilon \sim \mathcal{N}(0, I)\]The forward pass is implemented as follows. It returns the action distribution; the action is then sampled from it with rsample().
def forward(self, feat: torch.Tensor) -> TanhNormal | OneHotDist:
x = self.mlp(feat).float()
if self.discrete:
return OneHotDist(x)
mean, std = x.chunk(2, -1)
mean = self.mean_scale * torch.tanh(mean / self.mean_scale)
std = F.softplus(std + self.raw_init_std) + self.min_std
return TanhNormal(mean, std)
The actor creates imaginary trajectories that spawn from the real states predicted by the representation model. For each step, we predict the action given the current state, and then predict the next state using the transition model.
Note
States computed by the world model and fed into the transition model are provided in detach mode. This prevents the gradients of the transition model from propagating back into the representation model.
Value model (Critic)
With the imagined trajectories and their corresponding actions established, the value model comes into play to estimate expected future rewards:
\[v_\psi(s_\tau) \approx \mathbb{E}_{q(\cdot|s_\tau)} \left[ \sum_{t=\tau}^{\tau+H} \gamma^{t-\tau} r_t \right].\]The critic is implemented as a dense neural network that takes imaginary states and outputs the mean of a Gaussian distribution (with the standard deviation fixed to 1). Dreamer uses the mean of this distribution as the value estimate since, by definition, the state value is the expected return starting from that state.
To train the actor and critic, we must estimate the true state values along the imagined trajectories using the $\lambda$-return calculation, which relies on the reward and critic models.
Lambda-Return Calculation
\[V_R(s_\tau) = \mathbb{E}_{q_\theta, q_\phi} \left[ \sum_{n=\tau}^{t+H} r_n \right] \qquad (4)\] \[V_N^k(s_\tau) = \mathbb{E}_{q_\theta, q_\phi} \left[ \sum_{n=\tau}^{h-1} \gamma^{\,n-\tau} r_n + \gamma^{\,h-\tau} v_\psi(s_h) \right], \qquad h = \min(\tau + k,\; t + H) \qquad (5)\] \[V_\lambda(s_\tau) = (1-\lambda) \sum_{n=1}^{H-1} \lambda^{\,n-1} V_N^n(s_\tau) + \lambda^{\,H-1} V_N^H(s_\tau) \qquad (6)\]These equations are from the original Dreamer V1 paper, where γ is a fixed discount constant. Imagined trajectories start from a real state $s_t$ and are rolled out for H steps, with τ indexing the imagined steps from t to t+H.
The λ-return (Eq. 6) is a weighted combination of all k-step returns, from one step up to the imagination horizon H. Each k-step return (Eq. 5) sums the discounted rewards over the next k imagined steps and then bootstraps from the critic’s value estimate at the final state $s_h$, where $h = \min(\tau+k, t+H)$ truncates the rollout whenever τ+k would extend past the imagination horizon.
Rather than relying on a single one-step TD (temporal difference) target, the λ-return, as in TD(λ), interpolates between the one-step bootstrapped target and the longest available return, which runs to the horizon before bootstrapping. The hyperparameter $\lambda \in [0, 1]$ controls this interpolation by exponentially weighting the k-step returns. Small values of λ emphasize short-horizon bootstrapping, reducing variance but increasing bias; values closer to one emphasize longer returns, reducing bias at the cost of higher variance.
Starting with Dreamer V2, the fixed γ is paired with a learned continuation signal $c_t$, predicted by a continue model, giving an effective discount of $\gamma c_t$ that estimates whether the episode continues at each imagined step. λ thus determines how the k-step returns are blended, while the continue model accounts for possible early episode termination. In our implementation, the continue model is optional; when it is disabled, the fixed γ is used, as in Dreamer V1.
In the code, the return calculation is implemented as follows (the math derivation is in the Additional Material section at the end).
def lambda_return(
reward: torch.Tensor,
value: torch.Tensor,
discount: torch.Tensor,
lambda_: float,
) -> torch.Tensor:
"""TD(lambda) returns for an imagined trajectory of states s_0 .. s_H.
All inputs have shape (H + 1, ...), indexed by state. ``reward[t]`` and
``discount[t]`` belong to arriving at s_t, so the first entries are unused:
R_H = v(s_H)
R_t = r_{t+1} + discount_{t+1} * ((1 - lambda) * v(s_{t+1}) + lambda * R_{t+1})
Returns R_0 .. R_{H-1}, shape (H, ...).
"""
horizon = reward.shape[0] - 1
returns = []
last = value[-1]
for t in reversed(range(horizon)):
last = reward[t + 1] + discount[t + 1] * ((1 - lambda_) * value[t + 1] + lambda_ * last)
returns.append(last)
return torch.stack(returns[::-1])
Critic Loss
The value model (critic) $v_\psi(s_\tau)$ is trained by regression toward the λ-return $V_\lambda(s_\tau)$ from Eq. 6:
\[\min_\psi \; \mathbb{E}_{q_{\theta}, q_{\phi}} \left[ \sum_{\tau=t}^{t+H} \left| v_\psi(s_\tau) - V_\lambda(s_\tau) \right|^2 \right]\]The λ-return is a bootstrapped estimate, not ground truth: it combines predicted rewards with the critic’s own value predictions. Because it depends on the critic’s current parameters, gradients are stopped through it, so it’s treated as a fixed label for each update. This is standard in temporal-difference learning and makes the update a semi-gradient method
- Prediction: $v_\psi(s_\tau)$, the quantity being optimized.
- Target: $V_\lambda(s_\tau)$, held constant during the update (gradients stopped).
- Iteration: after the critic updates, its improved predictions produce improved λ-return targets, and the process repeats.
Intuitively, this works like a fixed-point iteration: the critic refines its estimates until they agree with the λ-return targets they help produce. What keeps this from being circular is the predicted rewards in the target, which are learned from real experience. Two caveats apply. First, with neural networks and a policy that changes during training, convergence isn’t guaranteed. Second, the values are consistent with the world model’s imagined dynamics, so their accuracy depends on how well that model matches the real environment. Dreamer V1 relies only on the stop-gradient for stability; Dreamer V2 added a slowly updated target critic for the bootstrap.
# excerpt from Dreamer.train_behavior
value_dist = self.value(feat[:-1].detach())
value_loss = -(weights * value_dist.log_prob(returns.detach())).mean()
Actor Loss
Once $\lambda$-returns are computed, they’re used to train the actor. The action model $q_\phi(a_\tau \mid s_\tau)$ is trained to produce actions that lead to trajectories with high value estimates:
\[\max_\phi ; \mathbb{E}_{q_{\theta}, q_{\phi}} \left[ \sum_{\tau=t}^{t+H} V_\lambda(s_\tau) \right]\]In practice, because optimizers minimize loss, this translates to minimizing the negative mean of the λ-values, each weighted by the probability that the imagined episode is still running at that step (the cumulative product of the discounts), with an optional entropy bonus:
# excerpt from Dreamer.train_behavior
actor_loss = -(weights * returns).mean()
if t.actor_entropy:
actor_loss = actor_loss - t.actor_entropy * (weights * entropy).mean()
Optimization Process
Putting these pieces together, a single optimization step unfolds across the imagined trajectory:
- Rollout Predictions: The reward model estimates rewards along the trajectory, while the critic predicts state values, paired with discounts: the fixed $\gamma$, multiplied by the predicted continuation probability at each step when the continue model is enabled.
- Target Generation: The λ-returns are computed from these predicted rewards, values, and discounts.
- Actor Update: The actor loss is evaluated as the negative discount-weighted mean of the λ-values to encourage high-value action policies.
- Critic Update: The critic loss is computed as the discount-weighted negative log-likelihood of the λ-returns under the predicted value distribution, with gradients explicitly detached on the target.
Environment Interaction
After each imagination-based training phase, the agent steps back into the real environment to evaluate its current policy and collect fresh experience for the replay buffer.
This follows a recurrent loop: a new observation arrives and is embedded by the encoder; the RSSM combines this embedding with the previous hidden state and action to compute the current stochastic state:
\[s_t \sim p_\theta(s_t \mid s_{t-1}, a_{t-1}, o_t)\]The actor samples an action from this state:
\[a_t \sim q_\phi(a_t \mid s_t)\]Exploration noise is added before the action is sent to the environment, so the agent keeps discovering new transitions instead of only exploiting what it already knows. The environment returns the next observation and reward; each step is stored in the replay buffer as the new observation together with the action that led to it, the reward received, and flags marking episode starts and true terminations. Once enough new experience accumulates, training resumes on the updated buffer.
Why Learning in Imagination Works?
One may wonder how is the learning in imagination even possible? Since if the world model is itself learned, isn’t training actor and critic on imagined trajectories circular?
It isn’t, because the signal is always grounded in real interaction. The dynamics model is trained via supervised learning on real transitions including decoding observations, predicting rewards, predicting episode continuation and this outer loop anchors everything else.
Even though a specific imagined trajectory was never actually visited, the world model has learned to condition transitions on actions and predict rewards from resulting states by generalizing from real transitions. So the reward signal used in imagination reflects structure learned from real experience.
Conclusion
If you’ve made it this far, well done: that was a lot to take in. With all the pieces in place, let’s step back and look at what the Dreamer framework has actually achieved.
Dreamer V1 showed that behavior learned purely in imagination could match leading model-free methods on continuous control while using far fewer real environment steps. That efficiency is a major advantage in domains like robotics, where real interaction is slow or costly. Dreamer V2 extended the approach to discrete, high-dimensional tasks and became the first world-model agent to reach human-level performance on Atari. Dreamer V3 went further still, using a single fixed set of hyperparameters across dozens of domains, including learning to collect diamonds in Minecraft from scratch.
Across all three versions, the core idea of training an actor and critic in imagination stayed the same. The gains came mainly from a stronger, more stable world model underneath it.
If you’d like to go deeper into Dreamer V1, the best next step is to see these ideas in code. Our implementation is available in this repository.
Additional Materials
Latent dynamics loss function derivation
\[\max I(s_{1:T}; (o_{1:T}, r_{1:T}) \mid a_{1:T}) - \beta I(s_{1:T}; i_{1:T} \mid a_{1:T}) \qquad (13)\]The first term is a conditional mutual information between our state sequence and joint observation-reward signal conditioned on action.
\[I(s_{1:T}; (o_{1:T}, r_{1:T}) \mid a_{1:T})\]By definition of conditional mutual information:
\[I(X; Y \vert{} A) = \mathbb{E}_{X,Y,A} \left[ \log \frac{p(Y \vert{} X, A)}{p(Y \vert{} A)} \right]\]Applying this definition to Dreamer V1 we get:
\[I(s_{1:T}; (o_{1:T}, r_{1:T}) \mid a_{1:T}) = \mathbb{E} \left[ \sum_t \log p(o_t, r_t \mid s_t, a_t) - \log p(o_t, r_t \mid a_t) \right]\] \[= \text{const} + \mathbb{E} \left[ \sum_t \log p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) \right]\]Where $\log p(o_t, r_t \mid a_t)$ is dropped because it is independent of the latent model’s parameters and acts only as an additive constant with respect to optimization. The exact likelihood $p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T})$ is intractable to compute directly. We introduce an approximate factorized decoder distribution:
\[q(o_{1:T}, r_{1:T} \mid s_{1:T}) = \prod_{t=1}^T q(o_t \mid s_t) q(r_t \mid s_t)\]Since Kullback-Leibler (KL) divergence is always non-negative ($\text{KL}(P \parallel Q) \ge 0$), subtracting a KL divergence between $p$ and $q$ yields a valid lower bound:
\[\mathbb{E}_p \left[ \ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) \right] \ge \mathbb{E}_p \left[ \ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) \right] - \text{KL} \left( p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) \, \Big\Vert{}\, \prod_{t=1}^T q(o_t \mid s_t) q(r_t \mid s_t) \right)\]By definition of KL divergence:
\[\text{KL}(p \parallel q) = \mathbb{E}_{p} \left[ \ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) - \ln \left( \prod_{t=1}^T q(o_t \mid s_t) q(r_t \mid s_t) \right) \right]\]Using the logarithm rule $\ln \left(\prod x_i\right) = \sum \ln x_i$:
\[\text{KL}(p \parallel q) = \mathbb{E}_{p} \left[ \ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) - \sum_{t=1}^T \left( \ln q(o_t \mid s_t) + \ln q(r_t \mid s_t) \right) \right]\]Now substitute this back into the lower bound expression:
\[\text{Lower Bound} = \mathbb{E}_p \left[ \ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) \right] - \mathbb{E}_p \left[ \ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T}) - \sum_{t=1}^T \Big( \ln q(o_t \mid s_t) + \ln q(r_t \mid s_t) \Big) \right]\]Notice that the $\ln p(o_{1:T}, r_{1:T} \mid s_{1:T}, a_{1:T})$ terms subtract out completely:
\[= \mathbb{E}_{p(o_{1:T}, r_{1:T}, s_{1:T}, a_{1:T})} \left[ \sum_{t=1}^T \left( \ln q(o_t \mid s_t) + \ln q(r_t \mid s_t) \right) \right]\]This demonstrates that maximizing the mutual information between the latent state sequence $s_{1:T}$ and environment trajectory $(o_{1:T}, r_{1:T})$ reduces to maximizing the expected log-reconstruction likelihood of observations and rewards given the latent states at each time step $t$.
For the second term $I(s_{1:T}; i_{1:T} \mid a_{1:T})$, we use non-negativity of KL divergence to obtain an upper bound:
\[I(s_{1:T}; i_{1:T} \mid a_{1:T}) = \mathbb{E} \left[ \sum_t \log p(s_t \mid s_{t-1}, a_{t-1}, i_t) - \log p(s_t \mid s_{t-1}, a_{t-1}) \right]\] \[= \mathbb{E} \left[ \sum_t \log p(s_t \mid s_{t-1}, a_{t-1}, o_t) - \log p(s_t \mid s_{t-1}, a_{t-1}) \right]\] \[\le \mathbb{E} \left[ \sum_t \log p(s_t \mid s_{t-1}, a_{t-1}, o_t) - \log q(s_t \mid s_{t-1}, a_{t-1}) \right]\] \[= \mathbb{E} \left[ \sum_t KL\big( p(s_t \mid s_{t-1}, a_{t-1}, o_t) \;\|\; q(s_t \mid s_{t-1}, a_{t-1}) \big) \right]\]The formula above can be derived as follows. By the definition of mutual information
\[I(X;Y\mid Z) = \mathbb{E}_{p(x,y,z)} \left[ \log p(x\mid y,z) - \log p(x\mid z) \right].\]Substituting $X=s_{1:T}$, $Y=i_{1:T}$, $Z=a_{1:T}$ gives
\[I(s_{1:T};i_{1:T}\mid a_{1:T}) = \mathbb{E}_{p(o_{1:T},r_{1:T},s_{1:T},a_{1:T},i_{1:T})} \left[ \log p(s_{1:T}\mid i_{1:T},a_{1:T}) - \log p(s_{1:T}\mid a_{1:T}) \right].\]The latent dynamics model factorizes over time,
\[p(s_{1:T}\mid i_{1:T},a_{1:T}) = \prod_{t=1}^{T} p(s_t\mid s_{t-1},a_{t-1},i_t),\]and similarly,
\[p(s_{1:T}\mid a_{1:T}) = \prod_{t=1}^{T} p(s_t\mid s_{t-1},a_{t-1}).\]Taking logarithms,
\[\log p(s_{1:T}\mid i_{1:T},a_{1:T}) = \sum_t \log p(s_t\mid s_{t-1},a_{t-1},i_t),\]and
\[\log p(s_{1:T}\mid a_{1:T}) = \sum_t \log p(s_t\mid s_{t-1},a_{t-1}).\]Therefore, we get
\[I(s_{1:T};i_{1:T}\mid a_{1:T}) = \mathbb{E} \left[ \sum_t \left( \log p(s_t\mid s_{t-1},a_{t-1},i_t) - \log p(s_t\mid s_{t-1},a_{t-1}) \right) \right].\]Using linearity of expectation,
\[= \sum_t \mathbb{E} \left[ \log p(s_t\mid s_{t-1},a_{t-1},i_t) - \log p(s_t\mid s_{t-1},a_{t-1}) \right].\]The Dreamer V1 paper assumes that the dataset index uniquely determines the observation,
\[p(o_t\mid i_t) = \delta(o_t-\bar{o}_t),\]where $\delta(\cdot)$ is the Dirac delta.
Therefore conditioning on $i_t$ is equivalent to conditioning on $o_t$,
\[p(s_t\mid s_{t-1},a_{t-1},i_t) = p(s_t\mid s_{t-1},a_{t-1},o_t).\]Thus,
\[I(s_{1:T};i_{1:T}\mid a_{1:T}) = \sum_t \mathbb{E} \left[ \log p(s_t\mid s_{t-1},a_{t-1},o_t) - \log p(s_t\mid s_{t-1},a_{t-1}) \right].\]The true predictive distribution $p(s_t\mid s_{t-1},a_{t-1})$ is generally intractable. Introduce a tractable approximation, $q(s_t\mid s_{t-1},a_{t-1}).$
Using the non-negativity of KL divergence,
\[KL \left( p(s_t\mid s_{t-1},a_{t-1}) \| q(s_t\mid s_{t-1},a_{t-1}) \right) \ge 0,\]we obtain
\[\mathbb{E}_{p} \left[ \log p(s_t\mid s_{t-1},a_{t-1}) \right] \ge \mathbb{E}_{p} \left[ \log q(s_t\mid s_{t-1},a_{t-1}) \right].\]Since this term is subtracted inside the mutual information,
\[-\mathbb{E}_{p}[\log p] \le -\mathbb{E}_{p}[\log q].\]Therefore,
\[I(s_{1:T};i_{1:T}\mid a_{1:T}) \le \sum_t \mathbb{E} \left[ \log p(s_t\mid s_{t-1},a_{t-1},o_t) - \log q(s_t\mid s_{t-1},a_{t-1}) \right].\]The expectation is over the joint distribution
\[p(o_{1:T},r_{1:T},s_{1:T},a_{1:T}).\]Using the chain rule,
\[p(o,a,s) = p(o,a)\, p(s\mid o,a),\]the expectation can be written as
\[\mathbb{E}_{p(o,a,s)} = \mathbb{E}_{p(o,a)} \left[ \mathbb{E}_{p(s\mid o,a)} \right].\]Therefore,
\[\sum_t \mathbb{E} \left[ \log p-\log q \right] = \sum_t \mathbb{E}_{p(o,a)} \left[ \mathbb{E}_{p(s_t\mid s_{t-1},a_{t-1},o_t)} \left[ \log p-\log q \right] \right].\]Recall the definition:
\[KL(p\|q) = \mathbb{E}_{p} \left[ \log p - \log q \right].\]The inner expectation is exactly
\[\mathbb{E}_{p(s_t\mid s_{t-1},a_{t-1},o_t)} \left[ \log p(s_t\mid s_{t-1},a_{t-1},o_t) - \log q(s_t\mid s_{t-1},a_{t-1}) \right].\]Therefore,
\[= KL \left( p(s_t\mid s_{t-1},a_{t-1},o_t) \| q(s_t\mid s_{t-1},a_{t-1}) \right).\]Finally,
\[I(s_{1:T};i_{1:T}\mid a_{1:T}) \le \sum_t \mathbb{E}_{p(o,a)} \left[ KL \left( p(s_t\mid s_{t-1},a_{t-1},o_t) \| q(s_t\mid s_{t-1},a_{t-1}) \right) \right].\]The information bottleneck term is upper bounded by the expected KL divergence,
\[\boxed{ I(s_{1:T};i_{1:T}\mid a_{1:T}) \le \sum_t \mathbb{E}_{p(o,a)} \left[ KL \left( p(s_t\mid s_{t-1},a_{t-1},o_t) \| q(s_t\mid s_{t-1},a_{t-1}) \right) \right]. }\]Since the original information bottleneck objective subtracts this mutual information,
\[I_{\text{predict}} - \beta I_{\text{compress}},\]replacing the compression term by its upper bound produces a valid lower bound on the overall objective. This expected KL divergence is precisely the KL regularizer optimized in Dreamer V1.
Derivation of the λ-return recursion
\[V_\lambda(s_\tau) = (1 - \lambda)\sum_{n=1}^{H} \lambda^{n-1} V^n(s_\tau) + \lambda^{H-1} V^H(s_\tau)\]The $\lambda$-return is a weighted average of all $n$-step returns:
\[G_t^\lambda = (1-\lambda)G_t^{(1)} + (1-\lambda)\lambda G_t^{(2)} + (1-\lambda)\lambda^2 G_t^{(3)} + \cdots\]where
\[G_t^{(n)} = \sum_{k=0}^{n-1}\gamma^k r_{t+k} + \gamma^n v_{t+n}\]The first few returns are:
\[G_t^{(1)} = r_t + \gamma v_{t+1}\] \[G_t^{(2)} = r_t + \gamma r_{t+1} + \gamma^2 v_{t+2}\] \[G_t^{(3)} = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \gamma^3 v_{t+3}\]Substitute them into the weighted sum:
\[\begin{aligned} G_t^\lambda ={}& (1-\lambda)(r_t + \gamma v_{t+1}) \\ & + (1-\lambda)\lambda (r_t + \gamma r_{t+1} + \gamma^2 v_{t+2}) \\ & + (1-\lambda)\lambda^2 (r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \gamma^3 v_{t+3}) + \cdots \end{aligned}\]Notice every single term begins with $r_t$. So we can factor it out:
\[\begin{aligned} G_t^\lambda ={}& r_t \left[ (1-\lambda) + (1-\lambda)\lambda + (1-\lambda)\lambda^2 + \cdots \right] \\ & + \gamma(\cdots) \end{aligned}\]The bracket is a geometric series:
\[(1-\lambda)\sum_{k=0}^{\infty}\lambda^k = (1-\lambda)\frac{1}{1-\lambda} = 1\]Therefore $G_t^\lambda = r_t + \gamma(\cdots)$
So the immediate reward is always taken completely. After removing $r_t$, the remaining terms inside the parenthesis are:
\[\begin{aligned} & (1-\lambda)v_{t+1} \\ & + (1-\lambda)\lambda (r_{t+1} + \gamma v_{t+2}) \\ & + (1-\lambda)\lambda^2 (r_{t+1} + \gamma r_{t+2} + \gamma^2 v_{t+3}) + \cdots \end{aligned}\]Separate the first term and factor out $\lambda$ from the rest:
\[(1-\lambda)v_{t+1} + \lambda \Big[ (1-\lambda)(r_{t+1} + \gamma v_{t+2}) + (1-\lambda)\lambda(r_{t+1} + \gamma r_{t+2} + \gamma^2 v_{t+3}) + \cdots \Big]\]Now look carefully at the expression inside the brackets.
The bracket is exactly $G_{t+1}^\lambda$ because
\[G_{t+1}^\lambda = (1-\lambda)G_{t+1}^{(1)} + (1-\lambda)\lambda G_{t+1}^{(2)} + \cdots\]and
\[G_{t+1}^{(1)} = r_{t+1} + \gamma v_{t+2}\] \[G_{t+1}^{(2)} = r_{t+1} + \gamma r_{t+2} + \gamma^2 v_{t+3}\]which matches the bracket exactly. Therefore:
\[(\cdots) = (1-\lambda)v_{t+1} + \lambda G_{t+1}^\lambda\]Substituting back gives the recursive identity:
\[\boxed{G_t^\lambda = r_t + \gamma \left( (1-\lambda)v_{t+1} + \lambda G_{t+1}^\lambda \right)}\]Code
Dreamer V1 PyTorch implementation: repo