Dreamer: World Models for RL

Why Does This Exist?

Model-free reinforcement learning algorithms such as PPO and Soft Actor-Critic achieve impressive asymptotic mastery, but require millions of physical interactions with the environment. In physical robotics, industrial control, and complex virtual worlds, collecting millions of real transitions is prohibitively expensive, slow, and dangerous.

Conversely, traditional model-based RL approaches that plan in observation space (such as pixel-level model predictive control) buckle under high computational costs and compounding pixel blurriness over extended horizons.

Dreamer (Hafner et al., 2019–2023) fundamentally solves this trade-off by decoupling data collection from policy optimization through a learned latent world model:

1. The world model continuously learns compact latent dynamics directly from historical sensory observations in an experience replay buffer.
2. The actor-critic agent learns behaviors entirely by "dreaming"—rolling out simulated trajectories inside the latent space of the world model without ever rendering pixels.
3. Because latent transitions are fast matrix multiplications on GPU hardware, the agent generates tens of thousands of imagined transitions per second, exploring complex behaviors with superhuman sample efficiency.

Think of It Like This

How It Actually Works

The Dreamer Architecture and Evolution

Dreamer operates through three concurrent, asynchronous learning stages:

1┌─────────────────────────────────────────────────────────────┐2│ 1. REAL ENVIRONMENT INTERACTION                             │3│    Agent executes policy π_ψ(a_t | z_t) in real world      │4│    Saves transitions (x_t, a_t, r_t, γ_t) to Replay Buffer   │5└──────────────────────────────┬──────────────────────────────┘6                               │ Real data7                               ▼8┌─────────────────────────────────────────────────────────────┐9│ 2. LEARN WORLD MODEL (RSSM)                                 │10│    Minimizes ELBO: Visual Recon + Reward Loss + KL Balance   │11│    Latent state z_t = (h_t, s_t)                            │12└──────────────────────────────┬──────────────────────────────┘13                               │ Seed states z_014                               ▼15┌─────────────────────────────────────────────────────────────┐16│ 3. LATENT IMAGINATION (ACTOR-CRITIC)                        │17│    Roll out policy π_ψ for H = 15 steps in latent space      │18│    Compute multi-step λ-returns V^λ(z_τ)                     │19│    Update Actor (analytic gradients) & Critic (two-hot)     │20└─────────────────────────────────────────────────────────────┘

The Evolutionary Milestones: V1, V2, and V3

1. DreamerV1 (2019):

   • Parameterized stochastic latents $s_t$ as continuous multivariate Gaussians $\mathcal{N}(\mu, \sigma^2 I)$.
   • Introduced analytic gradient backpropagation: because the transition dynamics and reward predictor are fully differentiable neural networks, the Actor $\pi_\psi(a \mid z)$ receives direct analytic gradients $\nabla_\psi \mathbb{E} \left[ \sum_{\tau} \hat{r}_\tau + \hat{V}(z_H) \right]$ backpropagated backward through time across the latent trajectory.
   • Mastered visual continuous control benchmarks (DeepMind Control Suite) directly from pixels.

2. DreamerV2 (2020):

   • Categorical Latents: Replaced continuous Gaussian latents with an array of 32 discrete categorical variables, each with 32 classes ($32 \times 32 = 1,024$ discrete choices). Backpropagation through discrete samples is accomplished using the straight-through estimator: $$z_{\text{st}} = z_{\text{one-hot}} + p - \text{stop\_gradient}(p)$$ Discrete categorical latents prevented latent representation collapse, accurately modeled multimodal environment branches, and achieved human-level performance across the 55 Atari 200M benchmark games from raw pixels.
   • KL Balancing: Addressed the issue of the prior dominating the posterior by splitting the KL loss into separate updates with distinct learning rates: $$\mathcal{L}_{\text{KL}} = \alpha_{\text{prior}} D_{\mathrm{KL}}(\text{sg}(q) \parallel p) + \alpha_{\text{post}} D_{\mathrm{KL}}(q \parallel \text{sg}(p))$$ Setting $\alpha_{\text{prior}} = 0.8$ and $\alpha_{\text{post}} = 0.2$ trains the prior toward the posterior four times faster than the posterior is pulled toward the prior.

3. DreamerV3 (2023):

   • Scale-Invariant World Modeling: Designed to run across wildly different domains (Atari, continuous control, DMLab, Minecraft) with fixed hyperparameters.
   • Symlog Transformation: Compresses wide dynamic ranges of rewards and observations using a symmetric logarithmic mapping: $$\text{symlog}(x) = \text{sign}(x) \ln(|x| + 1)$$ Its inverse, $\text{symexp}(y) = \text{sign}(y)(\exp(|y|) - 1)$, maps predictions back to linear scale.
   • Two-Hot Categorical Returns: Rather than predicting scalar values with mean-squared error, the Critic predicts probabilities over 255 discrete symlog-spaced bins, preventing gradient explosion in environments with sparse massive rewards.
   • DreamerV3 became the first algorithm to mine diamonds in Minecraft from scratch without human demonstration data.

Generalized $\lambda$-Returns for Latent Imagination

During latent imagination, the agent starts from an initial encoded latent state $z_0 = (h_0, s_0)$ and rolls out the actor policy $\pi_\psi$ for a fixed horizon $H = 15$:

To balance the low variance of the Critic $V_\xi(z)$ with the unbiased nature of long rollouts, Dreamer evaluates imagined trajectories using generalized $\lambda$-returns $V^\lambda(z_\tau)$.

The $n$-step bootstrapped return from state $z_t$ is:

The $\lambda$-return geometrically weights all intermediate $n$-step horizon estimates:

Worked numerical example

Let us compute the 2-step $\lambda$-return and symlog scaling for an imagined trajectory segment.

Trajectory Setup:

• Horizon sub-window: $H = 3$, discount factor $\gamma = 0.9$, decay weight $\lambda = 0.8$.
• Imagined reward sequence: $\hat{r}_1 = 1.0$, $\hat{r}_2 = 2.0$, $\hat{r}_3 = 0.5$.
• Critic value estimates at latent states: $V(z_1) = 3.0$, $V(z_2) = 2.5$, $V(z_3) = 1.0$.

Step 1: Compute $n$-step returns for starting state $z_1$

1. 1-step bootstrapped return ($n=1$):

   $$G_1^{(1)} = \hat{r}_1 + \gamma V(z_2) = 1.0 + 0.9(2.5) = 1.0 + 2.25 = 3.25$$

2. 2-step bootstrapped return ($n=2$):

   $$G_1^{(2)} = \hat{r}_1 + \gamma \hat{r}_2 + \gamma^2 V(z_3) = 1.0 + 0.9(2.0) + (0.9)^2(1.0)$$ $$G_1^{(2)} = 1.0 + 1.80 + 0.81(1.0) = 3.61$$

Step 2: Compute $\lambda$-return at $z_1$

Weighting the 1-step and 2-step returns via $\lambda = 0.8$:

Step 3: Symlog Transformation (DreamerV3)

The Critic targets this return using the symlog transform:

Evaluating the inverse to verify exact preservation:

Code

1import math2from dataclasses import dataclass3from typing import List, Tuple4
5
6def symlog(x: float) -> float:7    """DreamerV3 symlog transformation: compresses extreme return scales."""8    sign = 1.0 if x > 0 else (-1.0 if x < 0 else 0.0)9    return sign * math.log(abs(x) + 1.0)10
11
12def symexp(y: float) -> float:13    """Inverse symlog: maps compressed representation back to original scale."""14    sign = 1.0 if y > 0 else (-1.0 if y < 0 else 0.0)15    return sign * (math.exp(abs(y)) - 1.0)16
17
18@dataclass19class CategoricalLatent:20    """Represents a discrete categorical latent variable with straight-through estimator."""21
22    one_hot: List[float]23    probabilities: List[float]24
25
26class DreamerImaginationEngine:27    """Demonstrates Dreamer's latent imagination and return calculation mechanisms:28
29    1. Straight-through categorical sampling (DreamerV2/V3)30    2. Multi-step lambda-return calculation (TD-lambda)31    3. Scale-invariant symlog mapping (DreamerV3)32    """33
34    def __init__(self, gamma: float = 0.9, lam: float = 0.8) -> None:35        self.gamma = gamma36        self.lam = lam37
38    def sample_categorical_straight_through(39        self, logits: List[float]40    ) -> CategoricalLatent:41        """Samples a discrete category while preserving gradients via straight-through."""42        exp_logits = [math.exp(val) for val in logits]43        denom = sum(exp_logits)44        probs = [val / denom for val in exp_logits]45
46        # Discrete argmax one-hot selection47        best_idx = probs.index(max(probs))48        one_hot = [1.0 if idx == best_idx else 0.0 for idx in range(len(probs))]49
50        # In a differentiable framework: z_st = z_onehot + probs - probs.detach()51        return CategoricalLatent(one_hot=one_hot, probabilities=probs)52
53    def compute_two_step_lambda_return(54        self,55        r1: float,56        r2: float,57        v2: float,58        v3: float,59    ) -> Tuple[float, float, float]:60        """Calculates 1-step, 2-step, and composite lambda-return for state z_1."""61        # 1-step bootstrap62        g1 = r1 + self.gamma * v263
64        # 2-step bootstrap65        g2 = r1 + self.gamma * r2 + (self.gamma**2) * v366
67        # Lambda return blend68        v_lambda = (1.0 - self.lam) * g1 + self.lam * g269
70        return g1, g2, v_lambda71
72
73# Initialize engine with parameters matching worked numerical example74engine = DreamerImaginationEngine(gamma=0.9, lam=0.8)75
76# 1. Straight-Through Categorical Latent Demonstration77test_logits = [0.2, 1.8, 0.5, 0.1]78latent = engine.sample_categorical_straight_through(test_logits)79print(f"Sampled one-hot category: {latent.one_hot}")80# -> Sampled one-hot category: [0.0, 1.0, 0.0, 0.0]81
82print(83    f"Softmax probabilities: {[round(p, 3) for p in latent.probabilities]}"84)85# -> Softmax probabilities: [0.122, 0.603, 0.164, 0.11]86
87# 2. Lambda-Return Calculation88r_1, r_2 = 1.0, 2.089v_2, v_3 = 2.5, 1.090
91g_1, g_2, v_lambda_1 = engine.compute_two_step_lambda_return(92    r1=r_1, r2=r_2, v2=v_2, v3=v_393)94
95print(f"1-step return G_1^(1): {g_1:.2f}")96# -> 1-step return G_1^(1): 3.2597
98print(f"2-step return G_1^(2): {g_2:.2f}")99# -> 2-step return G_1^(2): 3.61100
101print(f"Lambda-return V^lambda(z_1): {v_lambda_1:.3f}")102# -> Lambda-return V^lambda(z_1): 3.538103
104# 3. Symlog Transformation105symlog_target = symlog(v_lambda_1)106print(f"Symlog return target: {symlog_target:.4f}")107# -> Symlog return target: 1.5125108
109restored_target = symexp(symlog_target)110print(f"Symexp restored target: {restored_target:.3f}")111# -> Symexp restored target: 3.538112
113# Verification assertions114assert round(g_1, 2) == 3.25115assert round(g_2, 2) == 3.61116assert round(v_lambda_1, 3) == 3.538117assert round(symlog_target, 4) == 1.5125118assert round(restored_target, 3) == 3.538

Watch Out For

The Quick Version

• Dreamer decouples reinforcement learning into two concurrent loops: learning an internal world model from replay data, and training an actor-critic agent entirely inside imagined latent rollouts.
• DreamerV2 replaced continuous Gaussian latents with arrays of categorical latents ($32 \times 32$) using straight-through estimators and KL balancing, preventing representation collapse.
• DreamerV3 introduced scale-invariant symlog transformations and two-hot return distributions, enabling identical hyperparameters across continuous control, Atari, and Minecraft.
• Policy optimization runs across $H = 15$ step latent rollouts evaluated via $\lambda$-returns, achieving 10,000+ simulated steps per second without decoding pixels.

What to Read Next

• Learning Latent Dynamics
• Model-Based Deep RL
• MuZero and AlphaZero