Mastering Diverse Domains through World Models

In plain words

Most reinforcement learning programs need someone to tune them for each new problem. A setting that works for one game often fails for another, so applying the method to a new task is slow and expensive. This paper presents DreamerV3, one method with one fixed set of settings that is meant to work across very different tasks: 2D and 3D games, robot-style control, board-like tasks and a sparse-reward open-world game.

The core idea is to give the agent an internal model of its environment. The agent watches what happens, learns to predict the next situation from its own actions, and then practises inside that prediction. A second part judges how good each predicted situation is, and a third part picks the next action. Because most of the practice happens in imagination, the agent needs fewer real interactions to improve.

The method also includes a set of small safeguards for the many different scales of rewards and predictions that show up across tasks. Those safeguards are what let one configuration work everywhere. The authors report that the method is the first to collect diamonds in Minecraft starting from nothing, without human demonstrations or a curriculum of easier stages.

TL;DR

DreamerV3 is a model-based reinforcement learning agent. It learns a recurrent world model from replayed experience, then trains an actor and a critic on imagined rollouts of 15 steps § Learning algorithm tab. 4. Its robustness techniques (symlog transforms, twohot value outputs, percentile return normalization and KL free bits) let every benchmark use the same hyperparameters tab. 4. Use it when you need a general agent without per-task tuning, and budget for roughly 0.1 to 9 A100 GPU days per benchmark tab. 2.

Problem

Reinforcement learning algorithms can be applied directly to tasks similar to the ones they were designed for, but moving to new domains requires substantial effort, expertise and compute for hyperparameter tuning § Introduction. The paper argues that this brittleness limits the use of RL on expensive models or tasks where tuning is prohibitive § Introduction.

Specific obstacles that the paper names:

• Magnitudes of rewards, returns and inputs differ by orders of magnitude across domains. Squared losses on large targets can diverge, while absolute or Huber losses stall learning, and running-statistics normalization makes the target non-stationary § Robust predictions.
• Entropy-regularized policy gradients need a scale that depends on the reward scale and frequency, so a fixed entropy weight does not transfer § Actor learning.
• Minecraft Diamond has sparse rewards, exploration difficulty, long horizons and procedurally generated worlds. Earlier approaches used human expert data or curricula § Introduction.

Method

Overview. Three networks are trained concurrently from replayed experience: a world model that predicts outcomes of actions, a critic that judges the value of each predicted outcome, and an actor that chooses actions § Learning algorithm. The agent acts by sampling from the actor, with no lookahead planning § Critic learning.

Steps.

1. Encode and roll out the world model. An encoder maps observations to stochastic latents z_t, and a sequence model keeps a recurrent state h_t. Together (h_t, z_t) form the model state, from which rewards, continuation flags and reconstructions are predicted § World model learning eq. 1 fig. 3.
2. Compute the world-model loss. Prediction loss covers the decoder, reward and continue predictors. Two KL terms, the dynamics loss and the representation loss, use stop-gradients on opposite sides and are each clipped below 1 nat § World model learning eq. 2 eq. 3.
3. Imagine trajectories. From every replayed start state, the world model and actor generate H = 15 imagined steps § Critic learning tab. 4 agent.py.
4. Compute λ-returns. Imagined rewards and critic values are combined with discount γ = 0.997 and λ = 0.95 into bootstrapped returns R^λ_t § Critic learning eq. 5.
5. Train the critic. The critic outputs a categorical distribution over 255 exponentially spaced bins, and its value is the expectation under that distribution. It is trained with cross-entropy against twohot-encoded returns, plus a regularizer toward an exponential moving average (EMA) of its own parameters § Critic learning § Robust predictions eq. 5 eq. 10 eq. 11 eq. 12 tab. 4.
6. Train the actor. REINFORCE on the return-minus-value advantage, divided by max(1, S), plus an entropy bonus with weight η = 3e-4 § Actor learning eq. 6.
7. Update the return scale S. S is the smoothed 5th-to-95th percentile range of the batch of λ-returns § Actor learning eq. 7.
8. Train on replay as well. The critic also regresses on replayed trajectories with loss scale 0.3, using the imagination returns at the start states as value annotations § Critic learning tab. 4.

# World-model loss for a batch of replayed sequences  [eq. 2] [eq. 3]
# Shapes: posterior and prior logits (B, T, S, C); rewards, continues (B, T)

def world_model_loss(obs, post_logits, prior_logits, rewards, continues):
    # Prediction loss: decoder (symlog squared), reward (twohot), continue
    # (logistic)  [§ World model learning] [eq. 3]
    prediction_loss = (
        decoder_loss(obs) +
        reward_twohot_loss(rewards) +
        continue_logistic_loss(continues))
    # Dynamics KL: prior pulled toward the stop-gradient posterior  [eq. 3]
    dynamics_kl = kl(stop_grad(post_logits), prior_logits)
    # Free bits: clip each time step's KL below 1 nat  [eq. 3]
    dynamics_loss = max(1.0, dynamics_kl)
    # Representation KL: posterior pulled toward the stop-gradient prior
    representation_kl = kl(post_logits, stop_grad(prior_logits))
    representation_loss = max(1.0, representation_kl)  # free bits  [eq. 3]
    # Weights beta_pred = 1, beta_dyn = 1, beta_rep = 0.1  [eq. 2] [tab. 4]
    return (
        1.0 * prediction_loss +
        1.0 * dynamics_loss +
        0.1 * representation_loss)

# Lambda-returns over imagined rewards and critic values  [eq. 5]
# rewards, continues: (H,); values: (H + 1,) predicted by the critic mean

def lambda_returns(rewards, continues, values, gamma, lam):
    # Bootstrap from the last value: R_H = v_H  [eq. 5]
    returns = [values[-1]]
    for t in reversed(range(H)):
        # Blend the next-step value with the recursive return  [eq. 5]
        target = (1 - lam) * values[t] + lam * returns[0]
        returns.insert(0, rewards[t] + gamma * continues[t] * target)
    return returns  # R^lambda_t for t = 0 .. H - 1

# Critic: categorical over 255 bins, trained on twohot targets  [eq. 10-12]
# bins = symexp(linspace(-20, 20, 255))  [§ Robust predictions]

def value_prediction(logits):
    # Read out the expectation under the bin distribution  [eq. 10]
    return softmax(logits) @ bins

def critic_loss(logits, lambda_returns, slow_mean):
    # Twohot target: mass on the two bins that bracket each return  [eq. 12]
    target = twohot(lambda_returns, bins)
    # Regularizer toward slow critic mean, weight 1  [§ Critic learning]
    slow_target = twohot(stop_grad(slow_mean), bins)
    # Cross-entropy with soft targets  [eq. 11]
    return cross_entropy(target, logits) + cross_entropy(slow_target, logits)

# Actor: REINFORCE on normalized advantages plus entropy  [eq. 6] [eq. 7]
# returns: R^lambda; values: critic baseline; S: current return scale

def actor_loss(returns, values, log_probs, entropies, scale):
    # Critic-relative advantage, scaled down only for large returns  [eq. 6]
    advantage = (returns - values) / max(1.0, scale)
    # Stop-gradient on the advantage; entropy weight eta = 3e-4  [eq. 6] [tab. 4]
    return -(stop_grad(advantage) * log_probs + 3e-4 * entropies)

# Return scale S: smoothed range between 5th and 95th percentile  [eq. 7]
# Percentiles are taken over the whole batch of lambda-returns

def update_return_scale(returns, state):
    # Batch percentiles of this step's returns  [§ Actor learning]
    low = percentile(returns, 5)
    high = percentile(returns, 95)
    # EMA, weight 0.99 on old value  [code: embodied/jax/utils.py#L90-L91]
    state.scale = 0.99 * state.scale + 0.01 * (high - low)
    return state.scale

Robust input and output transforms § Robust predictions eq. 8 eq. 9. Vector observations are passed through symlog(x) = sign(x) ln(|x| + 1) before the encoder, and vector reconstructions use a symlog squared loss. The inverse is symexp(x) = sign(x)(exp(|x|) - 1). Reward and critic outputs use the symexp twohot scheme: softmax logits over symmetric bins B = symexp([-20, ..., +20]), prediction equal to the probability-weighted bin average eq. 10, and a twohot target loss eq. 11.

Architecture app. Implementation app. Model sizes tab. 3 fig. 3. The encoder and decoder are convolutional for images (stride-2 convolutions down to 6×6 or 4×4) and 3-layer MLPs for vectors. The sequence model is a GRU with block-diagonal recurrent weights in 8 blocks. The actor and critic are 3-layer MLPs; the reward and continue predictors are 1-layer MLPs. The categorical latents are sampled with straight-through gradients § World model learning.

Model sizes tab. 3. The model dimension d sets the sizes: 12M (d = 256), 25M (384), 50M (512), 100M (768), 200M (1024) and 400M (1536). The recurrent units are 8d. The 200M size is the default for all benchmarks except the two control suites, which use 12M app. Benchmarks tab. 2.

Optimizer app. Implementation. Per-tensor adaptive gradient clipping (AGC) at 30% of the weight norm, with eps = 1e-3. Then LaProp, an RMS normalization followed by momentum, with eps = 1e-20. Learning rate 4e-5. No weight decay, annealing or dropout tab. 4.

Replay app. Implementation. A uniform replay buffer with an online queue. Each minibatch first takes non-overlapping online trajectories and then fills with uniformly sampled trajectories. Latent states are stored and written back during training. Replay ratio is the number of replayed time steps trained per environment step, so the number of gradient steps per environment step is replay ratio / (batch size × batch length × action repeat). Example from the paper: replay ratio 32 on Atari with action repeat 4 and batch shape 16 × 64 gives one gradient step every 128 environment steps app. Implementation.

Hyperparameters

Values from tab. 4, which says the same values are used across all benchmarks. The Code column is what the official config sets configs.yaml.

Settings the paper states outside the table, or that the code adds:

Baseline PPO tab. 1: Acme implementation with observation normalization, reward normalization, reward clipping at 10 standard deviations, epoch batch 64 × 256, 3 epochs, minibatch 8 × 256, policy trust region 0.2, no value trust region, advantage normalization, entropy 0.01, discount 0.997, GAE λ 0.95, learning rate 3 × 10^-4, gradient clipping 0.5, Adam epsilon 10^-5. The paper tuned epoch batch size, learning rate and entropy scale app. Baselines.

Benchmark protocols tab. 2:

Seeds: 5 per benchmark for Dreamer and PPO, except ProcGen (1 seed); 10 for BSuite and for Minecraft app. Benchmarks.

Minecraft environment app. Minecraft. Built on MineRL v0.4.4 (Minecraft 1.11.2) with a flat categorical action space. Episodes end on death or after 36000 steps (30 minutes at 20 Hz). Early termination on breaking diamond ore is removed. The jump action is held for 200 ms. Observations: a 64×64×3 first-person image, an inventory count vector, a max-count vector, a one-hot equipped item, and health, hunger and breath scalars. Rewards: +1 for each of 12 milestones, once per episode (log through diamond), plus −0.01 per lost heart and +0.01 per restored heart, with the heart terms not evaluated by the authors app. Minecraft.

Results

Author-reported. Setup: Dreamer with fixed hyperparameters, 200M model except as noted, one A100 GPU per agent. Baselines are the paper's fixed-hyperparameter PPO and each benchmark's published expert methods § Results app. Baselines. Numbers are as printed in the cited tables.

Notes on the table above:

• Atari100K baselines are reported in the paper's own table, which does not match standard protocols for all methods. EfficientZero reaches 190 (gamer score) and 109 (median) with a changed environment configuration, and the paper says that makes the comparison difficult § Results tab. 10.
• Proprioceptive control uses the 12M model. Visual control uses the 12M model as well tab. 2.
• The paper reports that Dreamer outperforms PPO across all domains and that robustness techniques matter, with the KL objective of the world model being the most important in its ablation § Results (ablation figure; no figure number is given in the text).
• Scaling: model sizes from 12M to 400M and replay ratios were studied on Crafter and a DMLab task. Larger models reached higher scores with less interaction; more gradient steps per environment step reduced the interactions needed § Results (Scaling properties).

Paper vs code

Checked against the official repository danijar/dreamerv3 at commit e01491fad6434b2245a3b8ca201dd7faedcc458c (MIT). The paper does not link the code directly; the link comes from the project page § Results.

• LaProp β2. Paper says LaProp uses its default parameters, β1 = 0.9 and β2 = 0.99 app. Implementation. Code sets β2 = 0.999 configs.yaml agent.py. Implement 0.999.
• Continue target. Paper says the continuation flags c_t are in {0, 1} § World model learning. Code trains the continue head on 1 − 1/333 (≈ 0.997) for non-terminal steps and 0 for terminal steps agent.py. Imagined returns use the predicted continuation as the discount agent.py. Replay returns use a fixed discount 1 − 1/333 agent.py.
• Loss weighting. Equations 5 and 6 are plain sums over t. Code weights each imagined step by the cumulative predicted continuation probability, both in the critic loss and in the actor loss agent.py agent.py agent.py. The paper does not mention this weighting.
• Reward head. Paper says the reward predictor is trained with the symlog squared loss § Learning algorithm, World model learning but later says reward and critic use symexp twohot § Robust predictions. Code uses symexp twohot for the reward head configs.yaml. Implement twohot for reward.
• Imagination horizon. Text says the critic's prediction horizon is T = 16; table says H = 15 § Critic learning tab. 4. Code uses imag_length: 15 and builds 16 feature states, H + 1, including the start state agent.py.
• Return scale. Paper defines S as the EMA of the percentile range eq. 7. Code smooths the 5th and 95th percentiles separately with rate 0.01 and takes the difference utils.py utils.py. Because the EMA is linear and debiasing is off configs.yaml, this matches the paper.
• Free bits. Paper clips each KL below 1 nat eq. 3. Code applies max(kl, free_nats) per step rssm.py. Consistent.
• Symexp twohot summation. Paper says positive and negative bins should be summed separately from small to large app. Implementation. Code sums the two halves symmetrically so predictions are exactly zero at initialization outs.py. Consistent in effect.
• Minecraft. The paper uses the block-breaking setting from prior work app. Minecraft. Code sets a break speed of 100 configs.yaml. The paper does not state this value.
• Code-only details. Policy std bounds 0.1 to 1.0 configs.yaml, 255 bins configs.yaml, and 32 × 64 categorical latents configs.yaml are not in the paper text.

Pitfalls

• Unnumbered sections. The paper's sections are unnumbered, so § citations on this card name the section or subsection title. Equation numbers were derived from the LaTeX source order (each \eq is one numbered equation, twelve in all). Check them against the compiled PDF before quoting them.
• Twohot numerics. Bins span symexp(±20), about ±4.9 × 10^8. Compute twohot targets and readouts in float32. The code casts logits to float32 even when compute is bfloat16 outs.py configs.yaml.
• Symmetric bin sum. Summing the probability-weighted bins left to right gives a nonzero prediction at initialization, which the paper warns about app. Implementation. Sum the two halves separately as the code does outs.py.
• Zero-initialized heads. Reward and critic output weights are zero so the agent does not hallucinate rewards or values at start app. Implementation configs.yaml configs.yaml.
• Return normalization floor. The scale is max(1, S), so returns with range below 1 are not rescaled. Under sparse rewards this keeps exploration strong § Actor learning. Do not normalize by standard deviation or by the batch min/max, which the paper says fail § Actor learning.
• EMA convention. The paper writes EMA(·, 0.99). Code gives weight 0.99 to the old value and 0.01 to the new one utils.py.
• Replay ratio. Gradient steps per environment step are replay ratio / (batch size × batch length × action repeat). Get the batch and action-repeat definitions right or the step budget will be off by a large factor app. Implementation.
• Critic on replay. Replay returns bootstrap from imagined returns at the start states, so the imagination pass must run before the replay critic loss agent.py.
• Cost. The default 200M model has 8192 recurrent units and 16 parallel environments. Minecraft used 8.9 A100 GPU days and ProcGen 16.1 tab. 2. Use the 12M model for control suites, which the paper says reach the same performance there app. Benchmarks.
• Minecraft environment changes. Reproducing the reported Minecraft result needs the paper's environment modifications: no early termination on diamond ore, the 200 ms jump hold and the 36000-step limit app. Minecraft. Without them the sparse-reward structure differs.
• DMLab claim vs table. Text says Dreamer exceeds IMPALA and R2D2+ at 1B steps § Results. The table lists R2D2+ only at 10B steps, where it scores 85.4 against Dreamer's 71.4. Treat the table as authoritative and read the 1B comparison as IMPALA only tab. 8.
• Atari100K comparability. EfficientZero changed the environment configuration; the paper says the comparison is hard § Results. Compare against the standard environment settings in tab. 10.
• Results are author-reported. No independent reproduction is cited on this card. The paper gives 5 seeds for most benchmarks, one for ProcGen and 10 for Minecraft and BSuite app. Benchmarks.

Lineage

• Builds on: DreamerV1 (continuous control) and DreamerV2 (Atari). The paper's version updates the robustness techniques, the architecture (block GRU, RMSNorm, SiLU), the optimizer (AGC plus LaProp) and the replay buffer app. Previous Dreamer generations. Neither predecessor has a card in this wiki.
• Compared against: PPO (fixed hyperparameters, Acme implementation); IMPALA, Rainbow and R2D2+ in DMLab; MuZero, IRIS, TWM, SPR, SimPLe and EfficientZero on Atari and Atari100K; PPG on ProcGen; D4PG, DMPO, MPO, DDPG, SAC, DrQ-v2 and CURL on control § Results app. Baselines.
• Superseded by: not stated on this card.