Skip to content

Dreamer

A from-scratch DreamerV3 (Hafner et al., "Mastering Diverse Domains through World Models", https://arxiv.org/abs/2301.04104) agent: an RSSM world model with categorical latents, trained jointly with an actor and critic on imagined rollouts. Implements the paper's five headline robustness techniques, cross-checked directly against the official implementation (github.com/danijar/dreamerv3, dreamerv3/rssm.py and embodied/jax/agent.py) rather than assumed from the paper text alone:

  1. Symlog inputs/reconstruction (rlax.signed_logp1/signed_expm1, used by the SymlogEncoder and TwoHotHead in .models).
  2. Categorical latents (num_latents independent categoricals of num_classes each, "stoch"/"classes" in the official implementation) with straight-through gradients and 1% "unimix" - mixing a little uniform mass into every categorical, so no class ever gets a literal zero probability - for both the prior and posterior.
  3. KL balancing with free bits: two separate KL terms with different stop-gradient placement and independent free-nats floors, not one combined KL clamped by a single scalar.
  4. Symexp-twohot regression for reward and value (TwoHotHead), not a Gaussian/MSE head - a discrete classification loss over an exponentially-spaced grid of bins, which is far less sensitive to reward-scale outliers than a Gaussian likelihood.
  5. Return normalization: an EMA-tracked 5th-95th percentile range of returns rescales advantages, so the policy gradient's magnitude stays stable across environments with very different reward scales, without per-environment tuning.

Also matches the official implementation's actor loss shape: REINFORCE (log_prob(action) * advantage + entropy bonus), not backpropagation through sampled discrete actions - discrete distrax.Categorical.sample() has no gradient path back to its logits, so naively differentiating an imagined-rollout return through a sampled discrete action (as an earlier draft of this agent did) trains nothing through that path at all - and an EMA "slow" target critic that the online critic is regularized toward, for training stability.

Deliberate simplifications, kept for navix's small grid-world observations rather than image-scale ones: a plain nn.GRUCell for the deterministic recurrent state (not the official "block GRU", a parameter-efficiency optimization for much larger deter sizes); a plain symlog+MSE decoder for observation reconstruction (not the official implementation's image-specific CNN decoder); ELU activations and no RMSNorm (not load-bearing for correctness, just a smaller/simpler net); num_latents/num_classes default to 8x8 rather than the paper's 32x32, sized for navix's small grids rather than Atari-scale observations. None of these are the algorithmic identity of DreamerV3 - the five techniques above are.

The world model's reusable building blocks (categorical-latent utilities, the symexp-twohot head, and the RSSM's encoder/decoder/prior/ posterior networks) live in .models, alongside PPO's shared network components; this module holds what's specific to Dreamer itself: the WorldModel that wires those blocks into an RSSM, the actor/critic heads, and the Dreamer agent's collection/replay/training loop.

Bases: Agent

Runs hparams.num_steps steps in hparams.num_envs parallel envs, carrying the per-env posterior latent (h, z, a_prev_oh) across steps so the policy always acts on an up-to-date belief.

Runs initialisation plus exactly one update() call. Mainly useful for tests/debugging, where running a full train() (sized by hparams.budget) is either unnecessary or awkward to size exactly.

Bases: HParams

Hyperparameters for Dreamer (DreamerV3). Frozen; .replace(...) for a variant. Groups the world-model, actor and critic knobs plus the imagination-rollout schedule; sized for navix's small observations rather than Atari.

Entropy bonus coefficient in the actor's REINFORCE loss.

Learning rate for the actor's optimizer.

Fraction of uniform probability mixed into the actor's action distribution wherever it's sampled from or scored (data collection, imagination, and the REINFORCE loss) - the same unimix_categorical technique already used for the world model's own categorical latents (see unimix_categorical's docstring), applied here too. Without this, the actor's entropy can reach exactly zero: once it does, collect_experience (which samples actions from this same distribution) stops exploring too, so real data collection narrows to whatever the collapsed policy repeats, the world model overfits to that narrow trajectory, and there is no path back - a self-reinforcing collapse verified empirically (entropy hits exactly 0.0 and success rate permanently flatlines at 0%, independent of how many actor gradient steps are taken per update - slowing that down only delays the same terminal collapse, it doesn't prevent it). A structural floor on the minimum action probability makes that specific failure mode impossible rather than merely less likely. Higher than the world model's 0.01 default since the action space here is much smaller (a handful of actions vs. many latent classes), so the same probability mass floor matters proportionally more per action.

Number of sequences sampled per world-model gradient step.

Number of bins for the reward/value symexp-twohot heads (paper: 255; reduced here since navix's rewards/returns span a far smaller dynamic range than Atari's).

Upper edge of the symlog-space bin range.

Lower edge of the symlog-space bin range.

Number of environment frames to train for.

Learning rate for the critic's optimizer.

Discount factor used in the imagined-rollout lambda-returns.

Weight of the KL(sg(post)||prior) ("dynamics") term.

Size of the encoder's output embedding.

Per-(batch,time) KL floor - below this, dyn/rep loss is zero.

Hidden layer size used throughout the model/actor/critic MLPs.

Length of the imagined rollouts used to train the actor and critic.

Lambda parameter of the imagined-rollout lambda-returns.

Flattened size of the categorical latent (num_latents * num_classes), i.e. how much of feat = concat([h, z_flat]) the latent occupies.

Maximum gradient norm for clipping, applied to each of the three optimizers independently.

Learning rate for the world model's optimizer.

Number of actor gradient steps per update (see num_model_updates for why the default is high).

Number of classes per categorical latent variable ("classes" in the official implementation).

Number of critic gradient steps per update (see num_model_updates for why the default is high).

Number of parallel environments to run.

Number of independent categorical latent variables ("stoch" in the official implementation).

Number of world-model gradient steps per update. The default is deliberately high relative to the frames collected per update (a replay ratio in the spirit of the official implementation's train_ratio): DreamerV3 is designed to be sample-efficient by gradient-stepping far more often than it collects. Concretely, with a sparse reward the rewarded transitions can be ~0.2% of the replay data, and a mean-reduced twohot cross-entropy pushes the reward head toward the base-rate prediction (~0) until it has seen enough positive examples to separate them - at 32 steps/update the reward head was still predicting ~0.004 at real goal transitions after a full 100k-frame run (policy stuck at random-walk success), while at 128 it reaches ~0.98 and the same run ends at 100% success.

Number of environment steps to collect per update.

Size of the RSSM's deterministic (GRU) hidden state ("deter" in the official implementation).

Weight of the KL(post||sg(prior)) ("representation") term.

Maximum number of environment frames kept in the replay buffer (rounded down to whole collection rollouts, and never allocated larger than the training budget itself can fill).

Floor on the return-normalization scale (perc95 - perc5), so a near-constant reward signal doesn't blow the advantage up by dividing by a near-zero scale.

EMA rate for the return-normalization percentile tracker.

Length of the sequences sampled for world-model training.

EMA rate the slow critic's params track the online critic at.

Weight of the online critic's regularization loss toward the slow critic's prediction, on top of its lambda-return regression loss.

Fraction of uniform probability mixed into every categorical.

Bases: PyTreeNode

Wraps three independent TrainStates, one per network - the world model, actor, and critic each have their own optimizer, learning rate and step counter, and are updated via their own apply_gradients(). An earlier draft instead subclassed TrainState directly and shared one tx/step field meant to be swapped between the three networks' updates; the swap never actually took effect before each network's update ran, so actor and critic gradients were silently applied through the model's optimizer (and learning rate) instead of their own, and the step counter never incremented since apply_gradients() was never called. Three separate TrainStates make both bugs structurally impossible instead of relying on remembering to swap a shared field correctly.

slow_critic_params is a plain EMA-tracked copy of the critic's params (not a fourth optimized TrainState - nothing ever computes a gradient w.r.t. it, it only ever gets updated by exponential averaging toward critic.params), used both to regularize the online critic's training and to compute the return-normalization statistics.

Builds an initial DreamerTrainState: inits the world model/ actor/critic's params and optimizers, resets hparams.num_envs environments, and allocates an empty replay buffer sized to whichever is smaller of hparams.replay_capacity and what hparams.budget can actually fill. Mirrors the flax.training. train_state.TrainState.create convention PPO's own TrainingState relies on (construction logic lives on the state itself, not on the agent that produces it).

Bases: PyTreeNode

A fixed-capacity FIFO replay buffer of whole collection rollouts ("blocks" of num_steps x num_envs transitions), carried inside DreamerTrainState so it lives through the jax.lax.scan training loop. DreamerV3 is a replay-based algorithm: the world model must keep training on past experience, not only the rollout just collected - with a sparse reward, a rare success transition seen only in the update it happened in is forgotten by the reward head one update later, so imagination goes back to predicting zero reward everywhere and the actor's learning signal vanishes as soon as it appears. (An earlier version had no replay at all - _sample_batch read directly from the latest rollout - which is exactly the failure mode that produced isolated windows of success that never consolidated.)

Blocks are written whole (one per Dreamer.update() call) at a rolling index; sampled sequences never straddle two blocks, so the rolling overwrite boundary can't splice unrelated timelines together. Only what world-model training needs is stored (obs flattened, action, reward, done, termination) - not the logging-only fields of Buffer.

Bases: Module

Rolls the RSSM forward purely through its own prior (no real observations), sampling actions from actor_logits_fn at each step - the "imagined" trajectories the actor and critic train on.

Returns (feats, rewards, terms, actions, action_logps), all (B, H, ...) (batch-major). rewards is the reward head's decoded mean (plain array, not a TwoHotHead logits tensor); actions/action_logps are returned because the actor loss needs log_prob(action) under the current actor params for its REINFORCE term - recomputed at loss-value-and-grad time from the stored actions, not reused from here (this rollout's actor calls may run in a different tracing context).

Touches every submodule exactly once, purely so WorldModel. init() can discover every parameter's shape. Parameters are keyed by submodule path, not by how many times a submodule is called in the "real" forward pass, so a single non-scanned pass creates the identical parameter tree observe() would - deliberately used instead of observe() for .init(), because .init() does its own internal tracing to discover shapes, which doesn't compose safely with tracing through observe()'s internal jax.lax.scan while already inside an outer jax.jit (surfaces as a confusing UnexpectedTracerError pointing at an unrelated submodule).

Runs the RSSM over an observed sequence, computing the posterior latent at every step and decoding obs/reward/term from it.

Parameters:

Name Type Description Default
obs_seq Array

f32[B, L+1, obs_dim], flattened observations.

 required
act_seq Array

i32[B, L], actions taken between consecutive observations.

 required
first_seq Array

f32[B, L], 1.0 where obs_seq[:, t + 1] is a fresh post-autoreset observation (the first of a new episode) rather than a real consequence of act_seq[:, t]. NOTE the required shift relative to the buffer's done flags: navix defers autoreset to the next env.step() call (Environment.step's should_reset looks at the INPUT timestep), so done[t] == 1 means obs[t + 1] is the genuine TERMINAL observation - the goal cell actually reached, caused by act_seq[:, t], carrying the episode's reward - and it's obs[t + 2] that is the exogenous reset. The correct mask for scan step t is therefore done[t - 1], which _sample_batch slices as a one-step-shifted window over done. An earlier version passed done UNshifted, masking one step early: it zeroed the belief state and action exactly at the goal-reaching transition, so the reward/term heads were trained to associate the sparse reward and termination with a blank-context posterior of the goal observation - a latent that imagination (which always rolls forward with full context) never produces, so imagined rollouts never saw reward at all, value targets stayed at zero, and the actor's advantage signal was identically ~0 despite the reward head fitting its (mislabeled) training data well. Meanwhile the actual garbage transition - the action consumed by the reset, "causing" the teleport to obs[t + 2] - was trained unmasked. Sampled training sequences are sliced from the replay with no regard for episode boundaries, so sequences commonly straddle an autoreset; at scan step t (which pairs act_seq[:, t] with embed(obs_seq[:, t + 1])), both the incoming (h, z) carried from step t - 1 and act_seq[:, t] itself are zeroed whenever first_seq[:, t] == 1, so the reset observation's posterior is computed from a blank slate instead of a stale, causally-unrelated belief plus an action that didn't really produce it - matching both the official implementation's is_first masking and what collect_experience does with its own carried latents (full-context belief at the terminal observation, blank slate at the reset one).

 required

Returns:

Type Description
Tuple[Array, Array]

`((h_seq, z_seq), (dyn_kls, rep_kls), feats, obs_pred, rew_logits,
Tuple[Array, Array]

term_logits), all aligned withobs_seq[:, 1:](L` steps,
Array

t=0..L-1) except (h_seq, z_seq) and feats, which keep
Array

their own (B, L, ...) shape. dyn_kls/rep_kls are
Array

f32[B, L], computed inside the scan body (not returned as
Array

distribution objects) because jax.lax.scan only restacks a
Tuple[Tuple[Array, Array], Tuple[Array, Array], Array, Array, Array, Array]

distrax.Distribution's array leaves across the new leading
Tuple[Tuple[Array, Array], Tuple[Array, Array], Array, Array, Array, Array]

axis; its batch_shape/event_shape (computed once from the
Tuple[Tuple[Array, Array], Tuple[Array, Array], Array, Array, Array, Array]

un-stacked per-step arrays) go stale, so calling e.g.
Tuple[Tuple[Array, Array], Tuple[Array, Array], Array, Array, Array, Array]

.kl_divergence() after the scan raises a shape mismatch deep
Tuple[Tuple[Array, Array], Tuple[Array, Array], Array, Array, Array, Array]

inside distrax/tfp.

One environment-collection step: advances the RSSM's belief with the previous action, then updates it to the posterior given the new observation. SymlogEncoder/RSSM/posterior are submodules bound to this WorldModel instance, so - unlike a bare Dense/Sequential module - they can only be called from inside a WorldModel.apply() trace, which is what this method (called via self.world.apply(..., method=WorldModel.posterior_step)) provides for collect_experience.

One imagined (prior/dynamics) transition with a caller-supplied action - the same mechanics as a single imagine rollout step, but decoupled from the actor. Returns (feat_next, reward_mean, term_prob). Used for diagnosis: comparing the reward/term heads' predictions on prior-sampled latents at known real transitions (e.g. forcing the recorded goal-reaching action from the recorded pre-goal posterior state) against their predictions on posterior latents isolates whether a learned signal actually survives into imagination, where the actor is trained.