Skip to content

DreamerV3

DreamerV3 learns a model of the environment and trains its policy inside that model. A recurrent world model learns, from replayed sequences, to predict the next latent state, the reward and whether the episode continues. An actor and a critic then learn mostly from trajectories imagined by the model. All three are trained together, one gradient step per update, with the LaProp optimizer and adaptive gradient clipping.

Paper: Hafner, Pasukonis, Ba and Lillicrap, Mastering Diverse Domains through World Models, arXiv:2301.04104v2, 2023; Nature, 2025 (arXiv:2301.04104v2). Official code: danijar/dreamerv3.

Ajax reproduces the code that produced the paper's numbers (commit 2411f7d) plus one later upstream bug fix (29eb964, the replay context). The world-model design notes give the full specification, and deviations.md lists every version choice and every departure from the reference.

Can I use?

Feature DreamerV3
Discrete actions yes
Continuous actions yes
Memory built in: the world model is recurrent (memory= is not accepted)
Environments gymnax, Brax, MuJoCo Playground
Observations flat vectors only; the environment must not normalise observations or rewards
Several seeds in one call yes, train(seed=[...])
Extension phases pretrain, post_update (after every update), eval_metrics
Resume and checkpoints yes; a resumed run continues the schedules and the replay, and equals the uninterrupted run (test)

Example

# doc: skip (compiles and trains for about a minute on a CPU; tests/agents/DreamerV3/test_dreamerv3_agent.py covers this path)
from ajax import DreamerV3
from ajax.logging.wandb_logging import LoggingConfig

# The smallest preset and a low train ratio, for a quick look on a CPU.
agent = DreamerV3("CartPole-v1", model_size="1m", n_envs=4, train_ratio=32,
                  batch_size=4, batch_length=16)
logging = LoggingConfig(config={}, log_frequency=1_000, use_wandb=False)
state, evaluations = agent.train(seed=0, n_timesteps=2_000, num_episode_test=2,
                                 logging_config=logging)
print(evaluations["Train/episodic mean reward"])  # the score the paper reports
print(evaluations["Eval/episodic mean reward"])

The paper's DMC-proprio protocol (Table 2): the 12M model, 16 environments, a train ratio of 512, action repeat 2, and 500K environment frames, which is 250K rows at repeat 2:

# doc: skip (the paper's full budget; benchmarks/world_models/paper_protocol.py runs it)
from ajax import DreamerV3

agent = DreamerV3("CheetahRun", model_size="12m", n_envs=16, train_ratio=512,
                  action_repeat=2, episode_length=1000)
state, _ = agent.train(seed=0, n_timesteps=250_000)  # 500K env frames

Key hyperparameters

Every default is the paper-era value.

Argument Default What it does
n_envs 16 environments stepped in parallel
model_size "12m" preset of "1m", "12m", "25m", "50m", "100m", "200m", "400m"; sets units, deter, hidden, classes
train_ratio 512 replayed steps trained per collected row
batch_size, batch_length 16, 64 sequences per batch, and their length
imag_horizon 15 imagined steps per start state
return_horizon 333 1 / (1 - gamma): sets the discount
lam 0.95 lambda of the imagined returns
learning_rate 4e-5 a float or a schedule of the update count; a 1,000-update linear warmup multiplies it
agc, eps 0.3, 1e-20 adaptive gradient clipping, LaProp's epsilon
actent 3e-4 actor entropy scale
free_nats 1.0 free bits of the KL losses
dyn_scale, rep_scale 1.0, 0.1 dynamics and representation loss scales
replay_capacity 5,000,000 rows kept, over all environments
action_repeat, episode_length 1, 1000 simulator steps per agent step, and per episode (Brax, Playground)

The other arguments (network depths, bins, unimix, the remaining loss scales, the slow critic and the return normaliser) are listed in the API below.

Notes

  • What n_timesteps counts. One row per environment per vector step, reset rows included, as the reference's step does. env_frames = rows * action_repeat is logged next to it.
  • The score. The paper reports the returns of the stochastic policy's training episodes, logged as Train/episodic mean reward. The evaluation also samples its actions, since the reference has no deterministic mode, and draws fresh keys at each evaluation (new initial states and action noise); a rerun of the same seed draws the same ones.
  • Actions. Continuous actions are clipped to [-1, 1] and mapped onto the environment's bounds, the reference's ClipAction and NormalizeAction.
  • Memory use. Each replay row stores the observation, the action, the reward, the flags and the posterior latent: about 8.3 KB at 12M on a 24-dimensional observation, so the 250K rows of the DMC protocol take about 2.1 GB per seed, and about 3.2 GB with the training step's working set. agent.replay_bytes_per_seed gives the figure after train. Running several seeds multiplies it.
  • Recompilation. Widths, depths, counts, bins, imag_horizon, batch_size, batch_length, train_ratio and replay_capacity fix the compiled program; changing one recompiles.
  • Main deviations (all in deviations.md): training is synchronous and in float32 (the reference is asynchronous and computes in bfloat16); one decoder head for the flat observation; everything is checkpointed and seeded. The two-hot expectation is computed in an exact form (D22); the reference's form adds noise to early predictions under JIT, which measurably changes the early actor entropy.
  • The DMC tasks run on MuJoCo Playground's MJX ports, whose physics differ from dm_control's (E18), so comparisons with the published curves are approximate.

Results

Recorded on CPU in PERFORMANCE_REPORT.md, with the 1M preset and 20,000 rows, before evaluations drew fresh keys: Pendulum swings up from 12,000 rows on (final evaluation -222, against about -1,200 at the start). On CartPole Ajax was compared with the reference code itself on the same protocol, over three pre-registered rounds; once both used the same two-hot expectation, the training statistics' seed ranges overlapped the reference's in at least 52 of 58 windows, and all five Ajax seeds passed the redefined CartPole check, as did all five of the reference's. The GPU validation against the published DMC curves is specified in VALIDATION.md and has not been run yet. The probes (tests/probing) train DreamerV3 briefly on tiny environments with known answers (discounting, a continuous advantage, and a risky-or-safe choice it can only weigh in imagination), and its own tests (tests/agents/DreamerV3) check that a run split by a checkpoint equals the run left whole. On the CPU speed benchmark (benchmarks/agent_bench.py: CartPole, 4 environments, 8,000 rows, a smaller model than the paper's) DreamerV3 ran at 367 rows per second of execution (275 with compilation) at main a8a13c2.

API

ajax.agents.DreamerV3.DreamerV3.DreamerV3 

DreamerV3(
    env_id,
    n_envs=16,
    model_size="12m",
    units=None,
    deter=None,
    hidden=None,
    classes=None,
    stoch=32,
    blocks=8,
    enc_layers=3,
    dec_layers=3,
    rew_layers=1,
    con_layers=1,
    actor_layers=3,
    critic_layers=3,
    bins=255,
    free_nats=1.0,
    unimix=0.01,
    actor_unimix=0.01,
    dyn_scale=1.0,
    rep_scale=0.1,
    rec_scale=1.0,
    rew_scale=1.0,
    con_scale=1.0,
    actor_scale=1.0,
    critic_scale=1.0,
    repval_scale=0.3,
    train_ratio=512,
    batch_size=16,
    batch_length=64,
    imag_horizon=15,
    return_horizon=333,
    lam=0.95,
    repval_lam=0.95,
    actent=0.0003,
    slowreg=1.0,
    slow_rate=0.02,
    retnorm_rate=0.01,
    retnorm_limit=1.0,
    minstd=0.1,
    maxstd=1.0,
    learning_rate=4e-05,
    agc=0.3,
    agc_pmin=0.001,
    beta1=0.9,
    beta2=0.999,
    eps=1e-20,
    warmup=1000,
    weight_decay=0.0,
    replay_capacity=5000000,
    episode_length=1000,
    action_repeat=1,
    env_params=None,
    extensions=(),
)

Bases: ActorCritic

DreamerV3, paper-era (2411f7d + 29eb964), for vector observations.

Discrete and continuous actions. n_timesteps (of :meth:train) counts rows, the reference's step: one per env per vector step, reset rows included (DESIGN.md section 2); env_frames = rows * action_repeat is logged next to it.

Paper DMC-proprio protocol (Table 2 p.19; dreamerv3_spec 7.5): the model size and train ratio of 2411f7d's dmc_proprio preset (configs.yaml:233-236), its action repeat (env.dmc.repeat: 2, :84) and Table 2's budget of 500K env frames, i.e. 250K rows at repeat 2 (the preset itself sets run.steps: 3e5, :237; dreamerv3_spec Open question 1)::

DreamerV3(
    task,
    model_size="12m",
    n_envs=16,
    train_ratio=512,
    action_repeat=2,
    episode_length=1000,
)
agent.train(seed=..., n_timesteps=250_000)  # 500K env frames

on the mujoco_playground ports of the DMC tasks (deviation E18), whose episodes are then 500 agent steps (1000 simulator steps). The score is Train/episodic mean reward, the returns of the stochastic policy's training episodes (dreamerv3_spec 7.3, 7.7).

Parameters:

Name Type Description Default
env_id Union[str, EnvType]

env id (gymnax, brax or mujoco_playground) or a prebuilt env. The env must not normalise observations or rewards.

 required
n_envs int

parallel envs (run.num_envs, 16).

 16
model_size str

preset '1m' (later code only, for small runs), '12m', '25m', '50m', '100m', '200m', '400m': model dimension d with units = hidden = d, deter = 8 d, classes = d / 16 (:data:~ajax.agents.DreamerV3.state.MODEL_SIZES).

 '12m'
units Optional[int]

explicit MLP width d, overriding the preset's.

 None
deter Optional[int]

explicit width of the deterministic state, overriding the preset's.

 None
hidden Optional[int]

explicit width of the RSSM's hidden layers, overriding the preset's.

 None
classes Optional[int]

explicit classes per latent, overriding the preset's.

 None
train_ratio float

replayed steps trained per collected row (an update every batch_size * batch_length / train_ratio rows after the training start; dreamerv3_spec 6.1-6.2).

 512
batch_size int

B, the windows of a batch.

 16
batch_length int

T, the trained rows of a window (each window has one more row, the replay context; spec 5.8).

 64
replay_capacity int

rows kept, over all envs (Table 4: 5e6); see "Replay memory" below.

 5000000
episode_length int

simulator steps per episode of brax / mujoco_playground envs (with action_repeat, episode_length // action_repeat agent steps); gymnax envs keep their own.

 1000
action_repeat int

simulator steps per agent step (brax / playground).

 1
env_params Optional[EnvParams]

gymnax env parameters.

 None
extensions Sequence[Extension]

Extensions folding the phases pretrain, post_update (after every update) and eval_metrics; others are rejected at construction.

 ()

Every other argument is the :class:~ajax.agents.DreamerV3.state.DreamerV3Config field of the same name and default, documented there.

Static hyperparameters (fixing shapes or the program, so changing them recompiles): every width, depth and count, bins, imag_horizon, batch_size, batch_length, train_ratio, replay_capacity; the scalar loss and optimizer coefficients are Python floats of the frozen configuration as well.

Replay memory. The replay is a ring of C = min(ceil( replay_capacity / n_envs), rows per env of the run) rows per env (capacity in rows: deviation D27), resolved by every :meth:train call (:meth:_resolve_replay: a run split into resumed calls keeps the replay of the uninterrupted run). Each row stores the observation, action, reward, flags and the posterior latent (deter float32 and stoch as uint8 class indices), about 4 (obs_dim + action_dim + deter + 1) + stoch + 3 bytes: 8.3 KB at 12m on a 24-dim observation, so the 250K rows of the DMC protocol take about 2.1 GB per seed (:attr:replay_bytes_per_seed gives the resolved figure). The training step needs a working set on top of it, mostly the imagination; it scales with B T (deter + stoch classes): about 1 GB at 12m (B = 16, T = 64, H = 15), so about 3.2 GB per seed for the DMC protocol in all (XLA's memory analysis of the compiled run). The seed vmap multiplies both.

replay_bytes_per_seed property 

replay_bytes_per_seed

Bytes of the resolved replay rings of one seed (0 before).

resume_iteration_offset 

resume_iteration_offset(initial_state)

The absolute tick of a resumed run, equal across seeds (:func:~ajax.environments.row_collector.resume_tick), so every schedule (the training-start gate, the ratio, the online queue, the logging cadence) continues instead of restarting.

train 

train(
    seed=42,
    n_timesteps=int(1000000.0),
    num_episode_test=10,
    logging_config=None,
    on_ids_ready=None,
    initial_state=None,
    total_timesteps=None,
    **kwargs,
)

Train for n_timesteps rows (n_timesteps // n_envs ticks).

Returns (state, evaluations) vmapped over seeds, as every agent on the shared loop (:mod:ajax.agents.loop): None without a logging config, else every logged key's values at the evaluations (every log_frequency rows). Resume with initial_state= a returned state (or (state, metrics)): the run continues the uninterrupted run of all the rows, schedules and replay included (:meth:_resolve_replay). total_timesteps is the run's horizon (:meth:ActorCritic.train).