Dreamer V2 — World Models for Model-Based RL

Use case: Autonomous vehicle control
Type: Model-based RL, latent dynamics, imagination-based actor-critic
Action space: Continuous

Overview

Dreamer V2 [Hafner et al., 2020] learns a compact latent world model from environment interactions, then trains an actor-critic policy entirely in imagination — inside the learned model — without additional real environment rollouts. This separation of world model learning from policy optimisation gives Dreamer exceptional data efficiency compared to model-free methods.

World Models: Core Idea

Instead of learning a policy directly from raw experience, Dreamer learns:

1. A world model that predicts latent dynamics, rewards, and episode termination
2. An actor-critic that is trained on imagined trajectories generated by the world model

The world model enables the agent to "think ahead" in latent space — thousands of imagination steps for every real step taken in the environment.

RSSM Architecture

The Recurrent State Space Model (RSSM) is the heart of the world model. It maintains a two-part latent state:

h_t  (deterministic): GRU hidden state — captures long-range temporal dependencies
z_t  (stochastic):    Gaussian latent  — captures the current scene uncertainty

The RSSM defines three distributions:

Transition prior:    p(z_t | h_t)            — imagination: predict next z from h
Posterior:           q(z_t | h_t, embed_t)   — training: infer z from real observation
Recurrent update:    h_t = GRU(h_{t-1}, z_{t-1}, a_{t-1})

The prior is used during imagination (no real observations). The posterior is used during world model training (conditions on encoded real observations). The two are trained to agree via a KL divergence loss.

Full Model Architecture

Latent Dynamics Learning

The world model is trained with four loss components:

1. KL Divergence Loss (with Free Nats)

Aligns the posterior (which sees real observations) with the prior (which does not):

KL_t = KL( q(z_t | h_t, embed_t) || p(z_t | h_t) )
L_KL = kl_scale * max( mean(KL_t) - free_nats, 0 )

Free nats: The free_nats threshold (default 3.0 nats) allows the model to use the posterior freely up to that information budget before being penalised. This prevents the prior from over-constraining the posterior early in training.

2. Reconstruction Loss

Ensures the latent state retains enough information to reconstruct observations:

L_recon = MSE( ObsDecoder(h_t, z_t), obs_t )

3. Reward Prediction Loss

L_reward = MSE( RewardPredictor(h_t, z_t), r_t )

4. Continue Prediction Loss

L_continue = BCE( ContinuePredictor(h_t, z_t), 1 - done_t )

Total world model loss:

L_WM = L_KL + L_recon + L_reward + L_continue

Imagination-Based Actor-Critic

After one world model update, Dreamer trains the actor and critic entirely in latent space:

1. Start from posterior states {(h_t, z_t)} collected during world model training
2. Unroll the RSSM for horizon steps using the current actor (no real environment needed):

   a_tau = actor(h_tau, z_tau)
   h_{tau+1} = GRU(h_tau, z_tau, a_tau)
   z_{tau+1} ~ prior(h_{tau+1})
   r_tau, done_tau = reward_predictor(h_tau, z_tau), continue_predictor(h_tau, z_tau)
   

3. Compute lambda returns from imagined rewards and critic values:

   G_t^lambda = r_t + gamma * cont_t * [(1-lambda)*V(h_{t+1},z_{t+1}) + lambda * G_{t+1}^lambda]
   

4. Actor loss: maximise lambda returns
5. Critic loss: MSE between critic predictions and lambda return targets

Lambda Returns

Lambda returns interpolate between 1-step TD (lambda=0) and Monte Carlo (lambda=1):

G_t^lambda = (1-lambda) * sum_{n=1}^{H-t} lambda^{n-1} * G_t^n  +  lambda^{H-t-1} * G_t^H

In practice, computed recursively backwards through the imagined trajectory:

G_H = r_H + gamma * cont_H * V_{H+1}
G_t = r_t + gamma * cont_t * ((1-lambda)*V_{t+1} + lambda*G_{t+1})

Comparison to Model-Free Approaches (Data Efficiency)

In practice, Dreamer achieves human-level performance on Atari with 200x fewer environment interactions than model-free methods.

Autonomous Vehicle Use Case

For the VehicleEnv, Dreamer learns a latent model of 2D vehicle dynamics, obstacle positions, and proximity to goal. After training the world model on ~500 real episodes, the actor-critic trains on imagined rollouts where the vehicle "mentally simulates" thousands of navigation trajectories per real step.

Key advantages over SAC in this setting:

• Sparse goal reward (only +10 on success) is easier to propagate through imagination
• The RSSM can learn to represent obstacle proximity implicitly in z_t
• The agent can reason about long-horizon planning (horizon=15) without excessive real rollouts

Key Hyperparameters

References

• Hafner, D., Lillicrap, T., Norouzi, M., & Ba, J. (2020). Mastering Atari with Discrete World Models (Dreamer V2). ICLR 2021. arXiv:2010.02193.
• Hafner, D. et al. (2019). Dream to Control: Learning Behaviors by Latent Imagination (Dreamer V1). ICLR 2020. arXiv:1912.01603.
• Hafner, D. et al. (2023). Mastering Diverse Domains through World Models (Dreamer V3). arXiv:2301.04104.
• Deng, F. et al. (2021). Latent Space Imagination for Efficient Reinforcement Learning.