DreamerV3 implementation reference: revised / Nature recipe, vector observations¶

This document merges three verified component specs (world model, actor-critic, training) into one reference. It includes every correction and missing item from the adversarial verifiers, and a re-check of the disputed points against the sources. It is meant to replace the paper and the code for the implementer.

Sources and pinned versions¶

Source	Identifier	Role
Paper	Hafner, Pasukonis, Ba, Lillicrap, "Mastering Diverse Domains through World Models", arXiv:2301.04104v2 (revised version, matches the Nature 2025 paper). https://arxiv.org/abs/2301.04104v2	Primary algorithm description. Page numbers below are v2 PDF pages.
Official code (pinned)	`danijar/dreamerv3` @ e3f0224 ("Fix Atari frame maxpooling on reset (#213)"). https://github.com/danijar/dreamerv3/tree/e3f02248693a79dc8b0ebd62c93683888ddaccfe	The "Spec (exact)" column describes this code unless marked otherwise. The world-model, actor-critic and config code at HEAD comes from the f8817c4 rewrite (2024-12-07). Later commits (55c8c67, ec19f16, cd026ac, bfcdfc1, b65cf81, cdf5709, e3f0224) touch only the API, logging, Docker/README, the atari100k task name, driver stop plumbing and Atari envs. None of them changes the algorithm.
Paper-era revised code	2411f7d (2024-04-15)	The code that produced the paper's Table 2 protocol. Used for version notes.
Post-paper fixes before f8817c4	1a532a7 (2024-04-28): removed the unused flags compute_lr, winit_scale and scale_by_actent, and renamed debug_outer to outer. 03bc307 (2024-05-08): one-hot encodes integer vector obs. 29eb964 (2024-05-15): fixes the first prevact in the replay context.	Version notes.
DreamerV3 v1 code	8fa35f8 (2023)	v1-vs-revised notes only.
External dependency	`elements` package (>= 3.19.1) provides `elements.when.Ratio` at HEAD. It is not vendored.	Its semantics are taken from `2411f7d:embodied/core/when.py:26-42`, the code it was moved from.

Reference sources: the paper PDF above and the official repository at the pinned commit. Code references path:line are relative to that repository.

Conventions¶

Dimensions. d = model dimension, which is also the MLP width ("units") and the RSSM "hidden" width. D = 8d is the deterministic state width, g = 8 is the number of GRU blocks, S = 32 is the number of categorical latents and C = d/16 is the number of classes per latent. B = 16 is the batch size, T = 64 the number of trained steps per sequence, K the number of imagination starts per row (K = T) and H = 15 the imagination horizon. N = B·K = 1024 is the number of imagined trajectories.
Recommended Ajax configuration. It uses d = 256 (12M preset): D = 2048, C = 16, feature width 10d = 2560. See the Hyperparameter defaults and Open questions sections.
Notation. sg = stop_gradient. symlog(x) = sign(x)·log(1 + abs(x)) and symexp(x) = sign(x)·(exp(abs(x)) − 1). Norms are written ‖·‖. Table cells avoid the pipe character, so absolute values are written abs(·).
Code references are relative to the repo root at e3f0224 unless prefixed with a commit, e.g. 2411f7d:dreamerv3/nets.py:321.
"Paper vs code" labels. agree; conflict (paper and code differ); code_only (the paper is silent); partial (some sub-points agree, others are code_only or conflict).
Version note. Gives differences at 2411f7d (paper era) and, where useful, v1 (8fa35f8). "Same" means algorithmically identical.
Where the paper and HEAD conflict, the "Spec (exact)" cell gives HEAD and flags the conflict. The recommended Ajax default is in the final table, and the choice is listed under Open questions.

Merge decisions and re-checks (traceability)¶

All verifier corrections are applied, and the verifiers' missing items are integrated:
sliding windows that cross episodes;
asynchronous write-back and parameter lag;
integer observations in the paper-era code;
weight-layout differences;
straight-through routing;
report windows;
eviction retry;
DMC preset history;
LR warmup labelled code-only;
RMSNorm precision history.
Re-checked in the source and decided:
replay_context = 0. The world-model spec said the official code "reuses the last carry of the previous unrelated batch, not a clean zero start". The training spec said it is a zero-state start. The source shows a zero start. Replay.sample always calls _annotate_batch(..., is_first=True), which sets is_first[:, 0] = True (embodied/core/replay.py:126, :284-286). RSSM._observe then masks deter, stoch and prevact to 0 wherever is_first holds (dreamerv3/rssm.py:75-80). The old carry is passed in but zeroed immediately, so the world-model spec's note is wrong. This item is corrected.
Asynchronous outs and params. The delayed write-back (pending_outs) and the acting-parameter lag of about two driver iterations (pending_sync holds the pre-update non-donated allo, swapped only after the next policy call) are confirmed at embodied/jax/agent.py:237-247, :264-289.
The first training batch is prefetched before the 1024-item gate. Confirmed: train.py:64 builds the stream at setup, Prefetch starts its worker in __iter__ (streams.py:44-48, :70-77), and Replay.sample waits only for len(sampler) >= 1 (replay.py:121-124).
No 1/2 factor in the symlog MSE in either code version. At 2411f7d TransformedMseDist.log_prob = −Σ(pred − symlog(y))², and at HEAD outs.MSE.loss = (pred − sg(symlog(y)))². The conflict with paper Eq. 8 holds for both versions.
Gradient-routing label. The paper's ablation (Ablations text p.10, Fig. 6b p.9, Fig. 18 caption: "additional reward and value gradients further improve performance") says that value gradients shape the representation by default. In the code, imagination features are stop-gradiented, so the replay critic (repval_grad=True) is the only path for value gradients into the world model. The training verifier's "agree (implicit)" therefore supersedes the world-model verifier's "code_only" for repval.
dyn-loss gradient reach (new correction found in the merge). The world-model spec said "dyn trains only the prior net". In the code, prior = _prior(feat['deter']) uses a deter that is not stop-gradiented (rssm.py:123), so dyn also trains the core and everything upstream of h_t. This matches paper p.4: "The dynamics loss trains the sequence model". Corrected in 2.14 and in the routing matrix (4.2).
Removed note: "AGC slows zero-initialised tensors". A constant per-tensor clip factor cancels in g / sqrt(EMA(g²)) when eps = 1e-20. At the first RMS step u = sign(g).
LaProp beta2. The paper (p.18) says beta2 = 0.99 and describes it as "its default parameters". Both code versions use 0.999. The defaults of the LaProp reference itself were not verified here (no literature fetch), so this stays open.

Version map: algorithmic differences, paper-era (2411f7d) to HEAD (e3f0224)¶

Each item is detailed in its own row below. This is only a summary.

Aspect	2411f7d (paper era)	HEAD e3f0224
dmc_proprio preset	size12m, train_ratio 512, env.dmc.repeat 2, run.steps 3e5 agent steps (6e5 env steps). Matches Table 2.	size1m (code only), train_ratio 1024, repeat 1, run.steps 1.1e6
Discrete actor distribution	'onehot' with 1% unimix (matches the paper)	'categorical' with no unimix
Start-state reward/continue in imagination	from replay data: reward_t and hard 1 − is_terminal_t	model predictions (r̂_0 unused, ĉ_0 ≈ 0.997 enters the weight)
Vector decoder output outscale	0.1 (Dist default; config outscale reached only imgout)	1.0
Symlog MSE	tolerance: errors² < 1e-8 zeroed; mode() returns symexp(pred)	no tolerance; pred() returns the symlog-space value
RMSNorm statistics	computed in the input dtype (bf16), upcast only for fp16	always computed in f32
Integer vector obs	cast to f32 and symlogged (2411f7d); then one-hot and symlog of everything (03bc307)	one-hot, not symlogged
First prevact in the replay context	bug: taken from the previous unrelated batch's carry (fixed by 29eb964 on 2024-05-15)	`data[action][:, K−1]`
Slow critic	separate module, f32, hard copy after the first optimizer step	copy at creation, compute dtype. Equivalent in effect.
Policy mean/std nonlinearities	computed after the cast to f32	tanh/sigmoid on bf16 Linear outputs, then cast to f32
Stored stoch	int32 argmax indices	f32 one-hot. Equivalent.
Encoder key order / decoder input order	obs-space order / concat(deter, stoch)	sorted keys / concat(stoch, deter). Layout only.
CNN (out of scope)	stride-2 convs, mults (1,2,3,4,4), debug_outer (first conv stride 1)	conv + 2×2 max-pool, mults (2,3,4,4)

1. Shared building blocks (used by world model, actor and critic)¶

These are the modular primitives worth implementing once and sharing. They are also reusable for TD-MPC2 where applicable: the symlog, two-hot and RMSNorm-MLP blocks.

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
1.1 Linear	`Linear(units)`. Kernel W has shape (in, prod(units)) and bias b has shape (prod(units),). The bias is always present (bias=True), even when RMSNorm follows, and is zero-initialised. y = x @ W.astype(x.dtype) + b.astype(x.dtype), reshaped to (..., units). `units` may be a tuple: heads emit (shape) or (shape, classes). The input must be in the compute dtype (bf16). Linear calls `ensure_dtypes(x)` but discards the return value, so only the forward* dtype assert runs. The custom_vjp backward (cotangent) assert is active only where the returned value is used, i.e. `Head.__call__`.	none	embodied/jax/nets.py:230-251; ensure_dtypes :102-117 (call sites :241, :268, :300, :375 discard the result); embodied/jax/heads.py:88	code_only	2411f7d: Linear also took norm/act kwargs and applied Linear → Norm → act internally (2411f7d nets.py:594-635). Same composition. v1 (8fa35f8 nets.py:570): no bias when a norm follows (`bias and norm == 'none'`).
1.2 BlockLinear (block-diagonal dense)	`BlockLinear(U, g)`. Input x has shape (..., I) with I % g == 0 and U % g == 0. Kernel shape (g, I/g, U/g); bias shape (U,), zero-initialised. y = einsum('...ki,kio->...ko', x.reshape(..., g, I/g), W).reshape(..., U) + b. Output block k depends only on input block k. Init fan_in = (I/g)·g = I, the TOTAL input width, because compute_fans treats the leading g axis like a conv receptive field. Std[W] ≈ 1/sqrt(I) is therefore sqrt(g) ≈ 2.83× smaller than a per-block fan-in would give. Only the forward dtype assert runs (see 1.1).	p.18 Networks: "GRU with block-diagonal recurrent weights of 8 blocks"	embodied/jax/nets.py:254-281; fans :187-197	agree on the block structure; fan convention code_only	2411f7d: block_fans=False (config default) gives the same total fan (2411f7d nets.py:873-884). It also padded I to a multiple of g, which never triggers for the RSSM.
1.3 RMSNorm, SiLU and the hidden block	RMSNorm over the last axis. x32 = f32(x); y = x32 · rsqrt(mean(x32², axis=−1, keepdims) + 1e-4) · scale. The learned scale (F,) is initialised to ones and stored in f32. There is no shift for 'rms'. The output is cast back to the input dtype (bf16). Activation: silu(x) = x·sigmoid(x). Every hidden layer is Linear → RMSNorm → SiLU (norm after the linear, before the activation, with the Linear bias kept). A LayerNorm ('layer') variant exists in the code but no config uses it.	Table 4 p.21 "Activation: RMSNorm + SiLU"; p.20 "Block GRU, RMSNorm normalization, SiLu activation"	embodied/jax/nets.py:361-409 (eps default :364; rms :382-386; scale init ones :401-405), :26-39 (act); dreamerv3/configs.yaml:91-101 (act silu, norm rms)	agree; eps code_only	2411f7d computed the RMS in the input dtype (bf16), upcasting only for float16 (2411f7d nets.py:753-758). HEAD always uses f32. Both use eps 1e-4. v1 used LayerNorm.
1.4 MLP and MLPHead	`nn.MLP(layers, units)` casts the input to bf16, flattens it to (N, F), applies `layers` × [Linear(units) → RMSNorm → SiLU] and reshapes back. `MLPHead(space, output, layers, units, outscale)` = MLP, then a `Head` (final Linear with outscale), then the output distribution. `outscale` reaches only the final head Linear; hidden layers always use outscale 1. `DictHead` builds one Head per key of a dict space.	p.18 Networks (MLP sizes per module)	embodied/jax/nets.py:565-587 (MLP); embodied/jax/heads.py:16-41 (MLPHead), :26-34 (outscale routed to Head only), :44-62 (DictHead), :65-94 (Head)	agree	2411f7d: same composition (Dist module per output).
1.5 Weight initialisation	Every Linear, BlockLinear and Conv kernel: W = outscale · 1.1368 · sqrt(1/fan_in) · TruncatedNormal(lower=−2, upper=2). The truncated normal has unit sigma before truncation and is sampled in f32. 1.1368 ≈ 1/0.8796 restores Std[W] ≈ 1/sqrt(fan_in). The config `winit: trunc_normal_in` parses to Initializer('trunc_normal', fan='in', scale=1.0). compute_fans: rank 0 gives (1, 1); rank 1 (n,) gives (1, n); rank 2 (in, out) gives (in, out); rank ≥ 3 gives (shape[−2]·prod(shape[:−2]), shape[−1]·prod(shape[:−2])). So a BlockLinear (g, I/g, U/g) gets fan_in = I, and a conv HWIO gets fan_in = H·W·I. Biases are zeros and never scaled by outscale. Norm scales are ones. All parameters are stored in f32. outscale per final layer: RSSM obslogit/priorlogit 1.0; vector decoder heads 1.0; reward head 0.0 (zero kernel and zero bias, so uniform softmax and a prediction of exactly 0); continue head 1.0; policy head 0.01 (both the mean and stddev Linears, or the logits Linear); value head 0.0. The slow value net starts as an exact copy of the value net (3.18). The RSSM initial state is zeros, not learned (2.2).	p.6 "initialize the output weight matrix of the reward predictor and critic to zeros"; p.18 "output weights of twohot distributions are initialized to zero". No general init scheme in the paper.	embodied/jax/nets.py:42-49 (name parsing), :144-197 (Initializer; trunc_normal :168-170; compute_fans :187-197), :250-251 (Linear._scaled_winit); embodied/jax/heads.py:26-34; dreamerv3/configs.yaml:91, 94-101	agree on zero twohot outputs; everything else code_only	2411f7d winit 'normal' is the same truncated normal × 1.1368 × sqrt(1/fan_in) with fan 'in' and block_fans False, so the fans are the same. Exception: the 2411f7d vector decoder output layer used outscale 0.1. Each vector head was `Dist(shape, dist=vecdist)` with no outscale passed, so the Dist default of 0.1 applied; the config outscale 1.0 reached only 'imgout' (2411f7d nets.py:321, :353, :408). HEAD passes outscale=1.0 to every vector head (rssm.py:304-305), so the HEAD decoder output kernel starts 10× larger. v1: fan 'avg'.
1.6 symlog / symexp	symlog(x) = sign(x)·log1p(abs(x)) and symexp(x) = sign(x)·(exp(abs(x)) − 1). Symlog of vector observations is computed in f32 (obs are float32) before the bf16 cast.	Eq. 9 p.7	embodied/jax/nets.py:59-64, :494-497	agree	Same.
1.7 DictConcat (input concatenation)	`DictConcat(spaces, fdims=1, squish)`: keys are sorted alphabetically. Availability mask per sample and per key (not per element): available(x) = (x != −inf for floats, x != −1 for signed ints) reduced with `.all()` over every feature dim. If any element of a key is unavailable, the whole vector for that key and sample is zeroed. The mask is applied before and again after encoding. Discrete keys become one_hot(int32(x), classes) in bf16 and are not squished. Continuous keys become squish(x) (symlog for the encoder, identity for actions), then bf16. Feature dims are flattened and the keys concatenated.	none	embodied/jax/nets.py:467-500 (sorted keys :471; masks :483-484, :497), :80-99 (available; .all() reduction :95-96)	code_only	2411f7d concatenated encoder vector keys in obs-space order (2411f7d nets.py:238-240). This changes only the layout.
1.8 Categorical / OneHot with unimix, straight-through, KL, entropy	Logits l of shape (..., S, C) are cast to f32. Unimix: p = 0.99·softmax(l) + 0.01/C, and l' = log(p). Sample (OneHot): idx ~ Categorical(l'); value = sg(one_hot(idx)) + p − sg(p), where p = softmax(l') is the mixed probability, so the straight-through gradient flows through the mixed softmax. Factorisation: `Agg(OneHot, dims=1, sum)` sums kl, entropy and logp over the S = 32 latent axis. KL(P‖Q) = Σ_S Σ_C softmax(l'_P)·(log_softmax(l'_P) − log_softmax(l'_Q)). Entropy = Σ_S −Σ_C p log p. Unimix is re-applied every time the distribution is built from logits, including after sg. Store raw logits. `outs.Categorical(logits)` without unimix (the HEAD discrete policy, 3.14) samples an int index with jax.random.categorical; logp = log_softmax(l)[a].	p.4 straight-through; p.5 "1% uniform and 99% neural network output"; p.18 Distributions; Table 4 "Latent unimix 1%"	embodied/jax/outs.py:208-240 (Categorical; unimix :210-217; kl :236-240; entropy :230-234), :243-270 (OneHot; straight-through :265-270), :40-76 (Agg)	agree	2411f7d: a tfp OneHotCategorical subclass with the same straight-through formula (2411f7d jaxutils.py:103-116), and Independent(..., 1) for the latent sum. Unimix was applied inside `_logit`, so the stored logit was already log(mixed probs). Mathematically identical.
1.9 Symexp two-hot distribution (reward and critic)	Bins (n = 255): half = symexp(linspace(−20, 0, 128, f32)); bins = concat([half, −reverse(half[:−1])]). This equals symexp(linspace(−20, 20, 255)) with exact antisymmetry, bins[127] = 0, a step of 20/127 in symlog space, and a range of ±symexp(20) ≈ ±4.85e8. Loss for target y (f32, raw space, not symlogged): see Algorithm D. The interpolation is linear in raw space between neighbouring symexp-spaced bins. A target outside the range puts all its mass on the edge bin; an exact bin hit puts all its mass on that bin. Prediction: probs = softmax(logits) in f32, then the symmetric sum: p[m]·b[m] + Σ(reverse(p[:m]·b[:m]) + p[m+1:]·b[m+1:]) with m = 127. This pairs mirror bins from small to large magnitude, so uniform probs give exactly 0. There is no unsquash.	Eq. 10-11 p.7 (B = symexp([−20 … +20])); Eq. 12 p.18 (twohot weights, k = Σ_j δ(b_j < x)); p.18 "summation order matters … positive and negative bins summed separately, from small to large"	embodied/jax/heads.py:132-144 (bins); embodied/jax/outs.py:273-330 (pred :285-309, loss :311-330); dreamerv3/configs.yaml:98, :101 (bins 255)	agree. The bin count is code_only. Eq. 12 and the code agree inside the range, including exact bin hits (Eq. 12 at x = b_{k+1} puts weight 1 on index k+1). The only difference is out-of-range targets: Eq. 12 is undefined for k = 0 or k = abs(B), while the code clips to the edge bin.	2411f7d emitted 256 logits and dropped the last for odd bins (2411f7d nets.py:432-438, :462-471). Same function class and the same symmetric mean (2411f7d jaxutils.py:210-250). v1 'symlog_disc': 255 bins linearly spaced in symlog space with target symlog(r), interpolation in symlog space and prediction symexp(E[bin]) (8fa35f8 nets.py:479-481, jaxutils.py:160-168).
1.10 Binary (Bernoulli) output	One scalar logit l, cast to f32. loss(c) = −[c·log_sigmoid(l) + (1 − c)·log_sigmoid(−l)], which accepts soft labels c ∈ [0, 1]. prob(1) = sigmoid(l). pred() = (l > 0).	p.4 "continue predictor via logistic regression"	embodied/jax/outs.py:189-205; embodied/jax/heads.py:96-99	agree	Same.
1.11 Symlog MSE output	pred = f32(Linear output). loss = (pred − sg(symlog(f32(target))))² elementwise, then summed over the key's feature dims by the Head's Agg. No 1/2 factor and no tolerance. pred() returns the symlog-space value (no symexp) and is used only in reports.	Eq. 8 p.7: L = ½(f(x) − symlog(y))²	embodied/jax/outs.py:129-141; embodied/jax/heads.py:85-94 (Agg sum over shape), :127-130	conflict: the code lacks the ½, so its reconstruction gradient is 2× the paper's formula (as if β_rec = 2)	2411f7d TransformedMseDist: also no ½ (log_prob = −Σ distance), but squared errors < 1e-8 were zeroed (tol) and mode() returned symexp(pred) (2411f7d jaxutils.py:179-207).
1.12 Mixed precision / dtypes	COMPUTE_DTYPE = bfloat16 (config `jax.compute_dtype`; internal.setup assigns `nets.COMPUTE_DTYPE`). f32: parameters, gradients (taken w.r.t. f32 params through the astype transpose), optimizer state (nu, mu), RMSNorm statistics, every output distribution's parameters (OneHot/Categorical logits, TwoHot logits and bins, Binary logit, Normal mean and std, MSE pred), and therefore all KLs, NLLs, entropies, log-probs, losses (the total is asserted to be an f32 scalar), λ-returns, normaliser statistics and Normal action samples. bf16: all matmul inputs and outputs (Linear casts the kernel and bias to the activation dtype), activations, the RSSM carry, deter, stoch and logits (rssm.py asserts COMPUTE_DTYPE), the GRU gate maths, one-hot action encodings, and encoder inputs after symlog. Sampled stoch is cast back to bf16. Latents returned to the host are upcast to f32 for replay storage. This is lossless: stoch is an exact one-hot and deter was bf16. Loss scaling runs only for float16 (4.9).	not specified	dreamerv3/configs.yaml:74; embodied/jax/internal.py:107-109; embodied/jax/nets.py:12, :18-23, :240-248, :374-399, :494-497; embodied/jax/outs.py:131-133, :163-165, :192, :211, :276; embodied/jax/opt.py:25-41; dreamerv3/rssm.py:62, :87, :91, :101-103; embodied/jax/agent.py:399-403	code_only	2411f7d: compute_dtype bfloat16 and param_dtype float32, the same policy. Differences: RMSNorm statistics were bf16 (1.3), the slow critic ran in f32, and policy tanh/sigmoid ran after the f32 cast (3.13). An all-f32 implementation changes numerics only, not the algorithm.

Algorithm D: two-hot loss and prediction (raw-space interpolation)¶

def twohot_loss(logits, y, bins):                      # logits [..., n] f32, y [...] f32 raw
    n = len(bins)
    below = clip(sum(bins <= y[..., None], -1) - 1, 0, n - 1)
    above = clip(n - sum(bins > y[..., None], -1), 0, n - 1)
    equal = below == above
    d_below = where(equal, 1, abs(bins[below] - y))
    d_above = where(equal, 1, abs(bins[above] - y))
    w_below = d_above / (d_below + d_above)
    w_above = d_below / (d_below + d_above)
    target = w_below[..., None] * one_hot(below, n) + w_above[..., None] * one_hot(above, n)
    return -sum(target * (logits - logsumexp(logits, -1, keepdims=True)), -1)

def twohot_pred(logits, bins):                         # odd n, m = (n - 1) // 2
    p = softmax(logits)
    m = (len(bins) - 1) // 2
    pb = p * bins
    return pb[..., m] + sum(pb[..., :m][..., ::-1] + pb[..., m + 1:], -1)

2. World model (RSSM, encoder, decoder, reward and continue heads, KL losses)¶

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
2.1 Latent state layout and size presets	Model dimension d. deter h_t has 8d floats (bf16 at run time). stoch z_t has S = 32 categorical latents × C = d/16 classes, stored as a one-hot of shape (32, C). The RSSM "hidden" width (width of the RSSM MLPs) is d. "units" (width of the encoder, decoder, reward, continue, policy and value MLPs; config regex `.\.units`) is d. The block-GRU has g = 8 blocks, each d wide. Feature for heads, actor and critic* = feat2tensor = concat([deter, stoch.reshape(32·C)]), width 8d + 2d = 10d. Fixed at every size: stoch 32, blocks 8, encoder 3 layers, decoder 3, reward head 1, continue head 1, policy 3, value 3, RSSM imglayers 2, obslayers 1, dynlayers 1, two-hot bins 255, RMSNorm + SiLU. LR, batch size and every other hyperparameter are identical across sizes. Presets as (d, deter, hidden, classes, CNN depth): size1m (64, 512, 64, 4, 4), code only; size12m (256, 2048, 256, 16, 16); size25m (384, 3072, 384, 24, 24); size50m (512, 4096, 512, 32, 32); size100m (768, 6144, 768, 48, 48); size200m = `defaults` (1024, 8192, 1024, 64, 64); size400m (1536, 12288, 1536, 96, 96). Feature width: 12M 2048 + 512 = 2560; 1M 512 + 128 = 640; 200M 8192 + 2048 = 10240. The class defaults in rssm.py (deter 4096, hidden 2048, classes 32, act gelu) are never used, because the config overrides them.	"Model sizes" and Table 3 p.20 (hidden d; recurrent units 8d; codes per latent d/16; "number of latents is fixed"); Table 2 p.19 (proprio and visual control use 12M; default 200M); Table 4 p.21 (B = 16, T = 64)	dreamerv3/configs.yaml:91 (rssm defaults: deter 8192, hidden 1024, stoch 32, classes 64, blocks 8), :94-101 (layers and units), :120-153 (presets), :178-182 (dmc_proprio); dreamerv3/rssm.py:16-32 (unused class defaults); dreamerv3/agent.py:51-53 (feat2tensor)	conflict: Table 3 prints 1024 recurrent units for 12M, which contradicts its own 8d rule (8 × 256 = 2048). This was verified on the PDF render of p.20. The code uses 2048. size1m is code only.	2411f7d: same presets (key `dyn.rssm`), no size1m. v1 (8fa35f8): default XL config with deter 4096, units 1024, classes 32, LayerNorm.
2.2 RSSM initial state and is_first reset semantics	initial(B) = {deter: zeros(B, 8d), stoch: zeros(B, 32, C)}, cast to bf16. It is not learned. At every observe step: (deter_{t−1}, stoch_{t−1}, prevact_t) ← where(is_first_t, 0, value) on the carry and the raw action. The action is then encoded (one-hot for discrete actions) and masked again with ~is_first, because one_hot(0) is not zero. The core then runs on these zeros, so at an episode start h_t = core(0, 0, 0). That is a constant set by the learned biases and norm scales, i.e. an implicitly learned initial state. The posterior then conditions on h_t and e_t as usual. imagine() never resets. Losses at is_first steps are not masked. For acting, the carry starts at zeros and the env's is_first resets it. For training, the carry comes from the replay context (5.8).	Not specified (Eq. 1 p.4 defines only h_t = f(h_{t−1}, z_{t−1}, a_{t−1}))	dreamerv3/rssm.py:45-49 (initial), :75-80 (mask, DictConcat, mask); embodied/jax/nets.py:67-77 (where/mask); dreamerv3/agent.py:101-107 (init_policy)	code_only	v1 (8fa35f8 nets.py:43-49, :89-93): learned initial state, deter = tanh(w) with w a learned vector initialised to zeros, and stoch = the mode of the prior from that deter. It was substituted at is_first via prev·(1 − is_first) + initial·is_first. Revised (2411f7d and HEAD): zeros plus masking.
2.3 Action encoding into the dynamics	`DictConcat(act_space, fdims=1)` (1.7). Keys sorted; availability mask per key and sample; discrete → one_hot(int32(x), classes) in bf16; continuous → identity (no symlog), cast to bf16. Inside `_core`: a ← a / sg(max(1, abs(a))) elementwise. Components with abs(a) > 1 become sign(a) in the forward pass, with gradient 1/max(1, abs(a)); one-hot actions are unaffected. Continuous policy samples (Normal(tanh(mean), std)) can leave [−1, 1]. Replay stores the UNCLIPPED policy sample (the ClipAction env wrapper clips only what the env sees), so this normalisation is what bounds the model's action input when acting, training and imagining. The driver zeroes the stored action at is_last steps. That is harmless, because the next step is is_first and is masked.	none	embodied/jax/nets.py:467-500, :80-99; dreamerv3/rssm.py:76-79, :97, :137; dreamerv3/main.py:249-258 (NormalizeAction + ClipAction); embodied/core/wrappers.py:76-86; embodied/core/driver.py:73-77	code_only	v1: the same operation (action_clip = 1.0, prev_action *= sg(1/max(1, abs(a))); 8fa35f8 nets.py:86-88, :106-108). 2411f7d is identical to HEAD (2411f7d nets.py:123).
2.4 Deterministic core: block GRU (sequence model f)	See Algorithm A. Code-only details that must be reproduced: (1) each input embedding is Linear + RMSNorm + SiLU, not purely linear; (2) there is one extra block-diagonal hidden layer `dynhid0` (BlockLinear → RMSNorm over the full D vector → SiLU) before the gate projection; (3) the gate projection `dyngru` has no norm and no activation; (4) the reset gate multiplies the whole candidate pre-activation (input and recurrent parts), unlike a textbook GRU that gates only the recurrent term; (5) u weights the candidate, and the −1 bias gives u ≈ 0.27 at init, so the cell keeps about 73% of the old state; (6) cross-block mixing comes only through x0, the embedding of the full deter, because each block's own deter slice and the gates are block-diagonal; (7) the per-block gate layout is [reset, cand, update] inside each block's 3D/g output chunk. Shapes at 12M: D = 2048, H = 256, per-block input D/g + 3H = 256 + 768 = 1024; dynhid kernel (8, 1024, 256) with fan_in 8192; dyngru kernel (8, 256, 768) with fan_in 2048. At 200M: per-block input 1024 + 3072 = 4096; dynhid (8, 4096, 1024); dyngru (8, 1024, 3072).	Eq. 1 p.4 (h_t = f_φ(h_{t−1}, z_{t−1}, a_{t−1})); p.18 Networks ("GRU with block-diagonal recurrent weights of 8 blocks … input to the GRU at each time step is a linear embedding of the sampled latent z_t, of the action a_t, and of the recurrent state to allow mixing between blocks"); p.20 "Block GRU"	dreamerv3/rssm.py:135-159	conflict on the embedding: the paper says linear, the code uses Linear + RMSNorm + SiLU plus the extra dynhid layer. The block-diagonal structure agrees.	2411f7d cell 'blockgru' is line-for-line the same (2411f7d nets.py:151-173), with block_norm False (one norm over the full D). v1 (8fa35f8 nets.py:103-140): dense GRU. img_in = Linear(concat(stoch, action)) + LayerNorm + SiLU, then one Linear(3·deter) on concat(deter, x), with no bias and LayerNorm applied to its 3·deter output (i.e. on the gate pre-activations, not on the input). The r/c/u formulas, including u = sigmoid(u − 1), are the same.
2.5 Prior / dynamics predictor p(ẑ_t ∣ h_t)	x = h_t (deter only). imglayers = 2 × [Linear(d) → RMSNorm → SiLU] (names prior0, prior1), then Linear(S·C, outscale 1.0) ('priorlogit'), reshaped to (S, C). The distribution is OneHot with unimix 0.01, factorised over S (1.8). In training, prior logits are computed in parallel over all (B, T) after the observe scan, from feat['deter'], the same h_t that conditioned the posterior. There is no separate rollout. In imagination they are computed inside the scan.	Eq. 1 p.4 "Dynamics predictor ẑ_t ~ p_φ(ẑ_t ∣ h_t)"; p.4 "dynamics, reward, and continue predictors are also MLPs"	dreamerv3/rssm.py:161-171 (_prior, _logit), :123 (prior from feat['deter'] in loss), :99 (in imagine); dreamerv3/configs.yaml:91 (imglayers 2, outscale 1.0)	agree. The layer count is code_only.	2411f7d: the same 2 img layers (2411f7d nets.py:112-117).
2.6 Posterior / representation q(z_t ∣ h_t, x_t)	x = concat([h_t, e_t], −1), where e_t are the encoder tokens (absolute=False; absolute=True would drop h_t). obslayers = 1 × [Linear(d) → RMSNorm → SiLU] ('obs0'), then Linear(S·C, outscale 1.0) ('obslogit'), reshaped to (S, C). z_t ~ OneHot(logits, unimix 0.01) is sampled with nj.seed() using straight-through (1.8), then cast to bf16. feat stores {deter, stoch, logit} with raw (pre-unimix) posterior logits. Stoch is sampled, not the mode, when acting, training and imagining.	Eq. 1 p.4 "Encoder z_t ~ q_φ(z_t ∣ h_t, x_t)"; p.4 straight-through	dreamerv3/rssm.py:81-92; dreamerv3/configs.yaml:91 (obslayers 1, absolute False)	agree	2411f7d is identical except that unimix was applied inside `_logit` (stored logits = log mixed probs; 2411f7d nets.py:66-70, :207-217). Mathematically the same.
2.7 observe loop (training and acting)	Training: tokens = enc(obs) for all (B, T) at once, then `nj.scan` over axis 1 (time) with unroll = 1. Carry {deter, stoch}; inputs (tokens_t, prevact_t, is_first_t). Each step runs Algorithm B. Outputs are the final carry, entries {deter, stoch} (B, T, ...) for the replay write-back (5.9), and feat {deter, stoch, logit} (B, T, ...). rssm.loss then computes prior = _prior(feat['deter']) for all (B, T). Acting: the same `_observe` with single=True on an (N_env,) batch. The carry (enc {}, dyn {deter, stoch}, dec {}, prevact) persists across env steps; prevact is the previously sampled raw action. `training=False` when acting, but the RSSM does not branch on it.	Eq. 1 p.4	dreamerv3/rssm.py:61-92, :120-133; dreamerv3/agent.py:115-135 (policy), :156-167 (loss)	agree	Same.
2.8 imagine loop (interface used by the actor-critic)	Per step (Algorithm C): a = policy(sg(carry)) (or the given action sequence when policy is not callable); deter = core(carry.deter, carry.stoch, DictConcat(a)); logits = _prior(deter); stoch ~ OneHot(logits, 0.01) with straight-through; return feat {deter, stoch, logit} and the action used. Scan over axis 1 for H steps. Imagined actions are int indices for discrete spaces and are one-hot encoded inside imagine. Only the policy input is stop-gradiented inside imagine; the core runs on non-sg start states. The world-model/actor-critic gradient cut happens in agent.py: imgfeat = concat([sg(first, skip=ac_grads), sg(imgfeat)], 1). `ac_grads` (False) controls only the sg on the start feat 'first'. Imagined feats are stop-gradiented unconditionally. With ac_grads = False, the world model gets no gradient from imagination. In training: starts = all B·T posterior entries (imag_last = 0 gives K = T), flattened batch-major to (B·K, ...), with H = 15.	p.5 Critic learning (imagination from replayed representations, T = 16); Table 4 (H = 15)	dreamerv3/rssm.py:94-118, :56-59 (starts); dreamerv3/agent.py:189-202 (sg at :196); dreamerv3/configs.yaml:88, :104-105	agree	2411f7d: only out['stoch'] was sg'd inside imgstep, and the start latent and action were sg'd (2411f7d agent.py:268-293). ac_grads was a string option 'none'/'first'/'all' (default 'none'), with the same effect.
2.9 Vector encoder	Keys: obs-space keys with len(shape) ≤ 2, excluding is_first, is_last, is_terminal and reward (and any 'log/' keys, removed in main.py). `DictConcat(vspace, 1, squish=symlog)` (1.7): keys sorted; continuous keys → symlog(x) in f32, then bf16; discrete (integer) keys → one-hot, not symlogged; unavailable keys zeroed per sample; feature dims flattened and concatenated, then reshaped to (B·T, F). Then layers = 3 × [Linear(d) → RMSNorm → SiLU] ('mlp0..2'). Tokens e_t = the output of the 3rd block, shape (B, T, d), with no final projection*. When image keys exist, CNN features are concatenated after the vector features.	p.18 Networks ("Vector inputs are symlog transformed and then encoded and decoded using 3-layer MLPs"); Eq. 9 p.7; p.5 symlog of vector observations	dreamerv3/rssm.py:179-224, :247-250; dreamerv3/agent.py:38-43; dreamerv3/main.py:130-131; embodied/jax/nets.py:59-64, :467-500; dreamerv3/configs.yaml:94 (layers 3, units, symlog True)	agree	Integer obs history: at 2411f7d itself, integer vector observations were cast to f32, concatenated and symlogged like continuous features, with no one-hot (2411f7d nets.py:249-253). 03bc307 added the one-hot but still symlogged the whole concatenated vector, one-hots included. HEAD one-hots and does not symlog. 2411f7d concatenated keys in obs-space order (HEAD sorts them). v1: 5 MLP layers with LayerNorm.
2.10 CNN encoder/decoder (out of scope for the vector-only build)	Image keys (uint8, rank 3) use a CNN encoder. Channels are (d/16)·mults with mults (2, 3, 4, 4), kernel 5, a stride-1 conv plus 2×2 max-pool per stage (strided False), and input /255 − 0.5. The decoder uses a block-linear spatial projection (bspace 8), nearest-upsample + conv, a sigmoid output, and per-pixel MSE summed over H, W, C. Not needed for vector observations.	p.18 Networks ("stride 2 convolutions … transposed stride 2 convolutions, with sigmoid activation"); p.20 (the layers closest to the data use d/16 channels)	dreamerv3/rssm.py:226-245, :308-356; dreamerv3/configs.yaml:94, :97	conflict (image only): HEAD uses conv + max-pool instead of strided convs, and its first conv has 2·d/16 = d/8 channels, against the paper's d/16.	2411f7d: kernel 5, mults (1, 2, 3, 4, 4), minres 4, and `debug_outer: True` for both enc and dec. The first encoder conv had stride 1 and the decoder's final imgout conv had stride 1; all other layers were stride 2 (transposed in the decoder) (2411f7d configs.yaml:117, :121; nets.py:263, :354).
2.11 Vector decoder (reconstruction loss)	Input: posterior feat as concat([stoch.reshape(32·C), deter]). Note the order: stoch first, unlike feat2tensor. Flattened to (B·T, 10d). `nn.MLP(3 layers, units d)` (Linear → RMSNorm → SiLU), then a `DictHead` with one Head per vector key, outscale 1.0. Continuous key ('symlog_mse', 1.11): pred = Linear(space.shape)(x), loss_key[b, t] = Σ over all of that key's feature dims of (pred − symlog(f32(target)))², with no ½ and no tolerance. Discrete key ('categorical'): logits = Linear(shape + (classes,)), loss = −log_softmax(logits)[target] summed over the shape dims. Targets are sg. Each obs key is its own loss term with scale rec = 1.0 (e.g. DMC walker: orientations, height, velocity are 3 terms). Reduction: sum over feature dims within a key, mean over (B, T), sum over keys.	Eq. 2-3 p.4-5 (L_pred = −ln p(x_t ∣ z_t, h_t) …); Eq. 8 p.7 (½ factor); p.18 Networks (3-layer MLP decoder); Table 4 (β_pred = 1)	dreamerv3/rssm.py:286-306 (input order :293); embodied/jax/heads.py:44-62, :65-94 (Agg sum :90-91), :101-110, :127-130; embodied/jax/outs.py:129-141; dreamerv3/agent.py:80-83 (rec scale per key), :178-182; dreamerv3/configs.yaml:86, :97	conflict: no ½ (1.11), so the reconstruction gradient is 2× the paper's formula	2411f7d: TransformedMseDist (tol 1e-8, mode = symexp), output outscale 0.1 (1.5), decoder input concat(deter, stoch) via Input(('deter', 'stoch')) (2411f7d nets.py:277). v1: 5-layer MLP with LayerNorm.
2.12 Reward head (symexp two-hot)	Input: feat2tensor(posterior feat) = concat([deter, stoch.flat]) (10d). It is not stop-gradiented (reward_grad = True), so reward gradients reach the RSSM and encoder (2.16). MLPHead(layers = 1, units = d): Linear(d) → RMSNorm → SiLU → Linear(255) with outscale 0.0 (zero kernel, zero bias). Output TwoHot(f32 logits, symexp bins) (1.9). Target: raw `obs['reward'][b, t]` in f32, no squash. Loss shape (B, T), scale 1.0.	Eq. 10-11 p.7; p.18 Networks (reward predictor is a 1-layer MLP); p.6 and p.18 (zero-init output weights)	dreamerv3/agent.py:55-57, :172-173; embodied/jax/heads.py:16-41, :132-144; dreamerv3/configs.yaml:98 (layers 1, outscale 0.0, bins 255), :114 (reward_grad True)	agree. The bin count is code_only.	2411f7d: 256 logits with the last dropped (same function class). v1: 5-layer head, symlog-space bins (see 1.9).
2.13 Continue head (Bernoulli) and its target	Input: feat2tensor(posterior feat), never stop-gradiented in the world-model loss. MLPHead(bool scalar space, 'binary', layers = 1, units = d, outscale 1.0): Linear(d) → RMSNorm → SiLU → Linear(units = ()) gives one scalar logit l. Target with contdisc = True: c_t = f32(~is_terminal_t)·(1 − 1/horizon) = 0.997·(1 − terminal_t), horizon = 333. Loss = soft-label BCE (1.10). Time-limit truncations (is_last without is_terminal) still get the target 0.997. Use is_terminal, never is_last. Imagination reads con.prob(1) = sigmoid(l) ≈ γ·P(continue) and uses an explicit disc = 1, because the discount is folded into the continue prediction (3.6).	p.4 ("continue predictor via logistic regression", c_t ∈ {0, 1}); p.5 (γ = 0.997); Table 4 (1/(1 − γ) = 333); p.18 (1-layer MLP)	dreamerv3/agent.py:56, :58, :174-177; embodied/jax/heads.py:96-99; embodied/jax/outs.py:189-205; dreamerv3/configs.yaml:99, :106-107 (horizon 333, contdisc True); dreamerv3/agent.py:401-405	conflict: the paper trains on binary c_t and applies γ in the λ-return (Eq. 5); the code trains on the soft label γ·(1 − term) and drops the explicit γ in imagination. The returns are equivalent in expectation, but the head's output semantics differ.	2411f7d is identical (contdisc True, soft label cont·(1 − 1/horizon); 2411f7d agent.py:258-261, :316-317). v1: 5-layer head, target 1 − terminal, discount applied separately.
2.14 KL losses (dynamics, representation) with free bits	post = feat['logit'] (B, T, 32, C), raw. prior = _prior(feat['deter']) (B, T, 32, C). dyn = KL(OneHot(sg(post), 0.01) ‖ OneHot(prior, 0.01)), summed over the 32 latents, giving (B, T). rep = KL(OneHot(post, 0.01) ‖ OneHot(sg(prior), 0.01)), giving (B, T). Both directions are KL(posterior ‖ prior). Unimix is applied after the sg. Free bits: dyn = max(dyn, 1.0) and rep = max(rep, 1.0), elementwise per (b, t) after summing over latents (not per latent). The gradient is zero where the KL is below 1 nat. jnp.maximum splits the gradient 0.5/0.5 on an exact tie, a measure-zero case. dyn and rep have identical forward values. Gradient reach (merge correction): prior = _prior(feat['deter']) is computed from a deter that is not stop-gradiented. So dyn (scale 1.0) trains the prior MLP and, through h_t, the core and everything upstream of h_t: earlier straight-through posterior samples, earlier encoder tokens, and action embeddings. rep (scale 0.1) trains the posterior obs layers and the encoder, and, through h_t, the core. The only stop-gradient is on the opposite distribution's logits. The extracted spec's claim that "dyn trains only the prior net" was wrong; paper p.4 says the dynamics loss "trains the sequence model". Computed at every (b, t), including is_first steps. Metrics: dyn_ent = mean prior entropy, rep_ent = mean posterior entropy (summed over latents).	Eq. 3 p.5; p.5 free bits ("clipping … below the value of 1 nat"); Table 4 (β_dyn = 1, β_rep = 0.1, free nats 1)	dreamerv3/rssm.py:120-133 (free_nats :127-129), :32; dreamerv3/configs.yaml:86 (dyn 1.0, rep 0.1), :91 (free_nats 1.0)	agree	2411f7d is identical (free 1.0, dyn 1.0, rep 0.1). v1 (8fa35f8 configs.yaml:72-74): β_dyn = 0.5, β_rep = 0.1, free 1.0.
2.15 Sequence data alignment	Replay step t holds x_t, r_t (the reward received on entering x_t; 0 at is_first, e.g. from_gym reset returns reward 0.0), is_first_t, is_last_t, is_terminal_t, and the action a_t chosen after observing x_t (zeroed by the driver if is_last_t). prevact_t = a_{t−1}, built as concat([carry_prevact[:, None], `data[action][:, :−1]`], 1), or from the replay context (5.8). Model: h_t = core(h_{t−1}, z_{t−1}, a_{t−1}); z_t ~ q(h_t, e_t). Heads at s_t = (h_t, z_t) predict x_t, r_t and c_t = (1 − terminal_t)·γ.	Eq. 1 p.4	dreamerv3/agent.py:316-318, :332, :173-177; embodied/envs/from_gym.py:61; embodied/core/driver.py:70-77	agree	Same.
2.16 Gradient routing into the world model	World-model params (enc, dyn, dec, rew, con) get gradients from: (a) the decoder losses; (b) rew (reward_grad = True, no sg on its input); (c) con; (d) dyn and rep; (e) the replay critic loss repval. For (e), repval_grad = True, so repl_loss runs the value net on feat2tensor(repfeat) without sg, and 0.3 × the critic regression backprops through the posterior features, the straight-through stoch, deter, the core and the encoder. Mechanism: every consumer of the posterior features (decoder, reward head, continue head, replay critic) receives z_t as a sampled one-hot carrying straight-through gradients, value = sg(onehot) + (p_mix − sg(p_mix)). These heads therefore backprop into the posterior logits (encoder and obs layers) as well as into deter. No gradient comes from imagination (2.8). Inside the world model, BPTT runs through the full T = 64 scan, via deter and via the straight-through stoch fed back into the core. The replay-context initial carry is data and is not differentiated.	p.4 ("world model parameters φ are optimized end-to-end"); Eq. 2-3 (L_pred includes −ln p(r_t ∣ …) and the continue term); p.4 straight-through; Ablations p.10, Fig. 6b, Fig. 18 caption (reward and value gradients shape the representation by default)	dreamerv3/agent.py:17 (sg with skip), :172 (reward_grad), :196 (ac_grads), :219-228 (repval_grad); dreamerv3/rssm.py:87, :96; embodied/jax/outs.py:243-270; dreamerv3/configs.yaml:88, :114-116	agree. Reward, continue and decoder routing are explicit in the paper. repval into the world model is implied by the ablation; its exact mechanism (replay critic with no sg, scale 0.3) is code detail.	2411f7d: the same (reward_grad True, replay_critic_grad True, ac_grads 'none'; 2411f7d agent.py:252, :353-354, :304-308).
2.17 Weight-layout differences (checkpoint compatibility only)	HEAD: encoder vector keys sorted alphabetically; decoder input concat(stoch, deter); heads, actor and critic take concat(deter, stoch); per-block GRU gate layout [r, c, u] (2.4). These matter only when loading official weights. They are mathematically equivalent function classes.	none	embodied/jax/nets.py:471; dreamerv3/rssm.py:293, :155-156	code_only	2411f7d: decoder input concat(deter, stoch) (2411f7d nets.py:277); encoder keys in obs-space order (2411f7d nets.py:238-240).

Algorithm A: block-GRU core (`RSSM._core`, dynlayers = 1)¶

def core(deter, stoch, action):     # deter [B, D] bf16, stoch [B, S, C] bf16 one-hot, action [B, A] bf16
    D, Hd, g = 8 * d, d, 8          # Hd = rssm.hidden
    z = stoch.reshape(B, S * C)
    a = action / sg(maximum(1, abs(action)))
    x0 = silu(rms(Linear(Hd)(deter)))          # 'dynin0'  (full deter: the only cross-block mixing)
    x1 = silu(rms(Linear(Hd)(z)))              # 'dynin1'
    x2 = silu(rms(Linear(Hd)(a)))              # 'dynin2'
    x = concat([x0, x1, x2], -1)               # [B, 3Hd]
    x = repeat(x[:, None, :], g, axis=1)       # [B, g, 3Hd]: same embedding for every block
    x = concat([deter.reshape(B, g, D // g), x], -1).reshape(B, g * (D // g + 3 * Hd))
    x = silu(rms(BlockLinear(D, g)(x)))        # 'dynhid0'; RMSNorm over the FULL D vector
    x = BlockLinear(3 * D, g)(x)               # 'dyngru'; no norm, no activation
    r, c, u = split(x.reshape(B, g, 3 * D // g), 3, axis=-1)   # per-block [r | c | u]
    r, c, u = (y.reshape(B, D) for y in (r, c, u))
    r = sigmoid(r)
    c = tanh(r * c)                            # reset gate scales the whole candidate pre-activation
    u = sigmoid(u - 1)                         # u weights the candidate; ~0.27 at init
    return u * c + (1 - u) * deter

Algorithm B: one observe step (`RSSM._observe`), used both when training and when acting¶

def observe_step(carry, token_t, prevact_t, is_first_t):
    deter, stoch, prevact_t = where(is_first_t, 0, (carry.deter, carry.stoch, prevact_t))
    act = DictConcat(act_space)(prevact_t)     # one-hot for discrete, identity for continuous
    act = where(is_first_t, 0, act)            # mask again: one_hot(0) != 0
    deter = core(deter, stoch, act)            # h_t
    x = concat([deter, token_t], -1)
    x = silu(rms(Linear(d)(x)))                # 'obs0'
    logit = Linear(S * C, outscale=1.0)(x).reshape(S, C)   # 'obslogit', raw
    stoch = bf16(OneHot(logit, unimix=0.01).sample())      # straight-through sample
    return {deter, stoch}, entry={deter, stoch}, feat={deter, stoch, logit}

Algorithm C: one imagine step (`RSSM.imagine`, single=True)¶

def imagine_step(carry, policy):
    a = policy(sg(carry))                      # pi(. | feat2tensor(sg(carry))).sample()
    deter = core(carry.deter, carry.stoch, DictConcat(act_space)(a))
    x = deter
    for i in range(2):
        x = silu(rms(Linear(d)(x)))            # 'prior0', 'prior1'
    logit = Linear(S * C, outscale=1.0)(x).reshape(S, C)   # 'priorlogit'
    stoch = bf16(OneHot(logit, unimix=0.01).sample())
    return {deter, stoch}, (feat={deter, stoch, logit}, action=a)

Algorithm E: world-model part of the loss (`Agent.loss`, first half)¶

tokens = enc(obs)                                          # [B, T, d]
carry, entries, feat = rssm.observe(carry0, tokens, prevact, obs.is_first)   # scan, Algorithm B
prior = rssm.prior(feat.deter)                             # [B, T, S, C], parallel
dyn = max(KL(OneHot(sg(feat.logit)) || OneHot(prior)), 1.0)       # summed over S
rep = max(KL(OneHot(feat.logit) || OneHot(sg(prior))), 1.0)
recons = dec(concat(stoch_flat, deter))                    # one dist per obs key
losses[key] = recons[key].loss(sg(obs[key]))  for each vector key   # [B, T] each
x = feat2tensor(feat)                                      # concat(deter, stoch_flat), NOT sg
losses.rew = rewhead(x).loss(obs.reward)                   # two-hot NLL, raw reward
losses.con = conhead(x).loss(f32(~obs.is_terminal) * (1 - 1/333))
losses.dyn, losses.rep = dyn, rep                          # all [B, T]

3. Actor-critic learning in imagination¶

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
3.1 Imagination start states	Imagination starts from every posterior state in the replayed batch. K = min(imag_last or T, T), and imag_last = 0, so K = T. With B = 16 and T = 64 that is B·T = 1024 start states. There is no subsampling. The replay sequence has 64 + 1 steps; the context step is dropped (5.8), leaving T = 64. Each start carry is the posterior entry {deter_t, stoch_t}, where stoch_t is the sampled one-hot posterior. Flatten batch-major: x[:, −K:].reshape(B·K, ...). The same ordering is used for the 'first' features and for reshaping the replay bootstrap to (B, K) (3.19). No start state is masked: terminal and is_last starts are still imagined, and are down-weighted only through the continuation weight (3.6).	p.5 Critic learning ("Starting from representations of replayed inputs…"); Table 4 (B = 16, T = 64)	dreamerv3/agent.py:189-191, :194-195; dreamerv3/rssm.py:56-59; dreamerv3/configs.yaml:10-11, :15, :104; dreamerv3/main.py:184-188	agree. Using all B·T states is code detail.	2411f7d: imag_start 'all' with imag_repeat 1 (2411f7d configs.yaml:132-133; agent.py:276-282), giving the same start set. batch_length was 65 including context, so T is also 64.
3.2 Horizon and trajectory layout	H = imag_length = 15 imagined steps. Each trajectory has H + 1 = 16 states: index 0 is the posterior start state and indices 1..15 are imagined prior states. The paper's "prediction horizon T = 16" counts these 16 states. imgfeat[:, 0] = sg(posterior feat at t); imgfeat[:, i] = sg(state after i imagined steps). imgact[:, i] = the action sampled at state i, for i = 0..H, so imgact is aligned with imgfeat. imgact[:, H] ('lastact') is sampled at the last state only to fill the shape, and is never used because every loss slices [:, :−1]. Shapes are (B·K, H + 1, ...). Losses cover t = 0..H−1 (15 steps). State H is used only for bootstrapping.	p.5 "prediction horizon T = 16"; Table 4 (H = 15)	dreamerv3/agent.py:190, :193-201; dreamerv3/rssm.py:94-118; dreamerv3/configs.yaml:105	agree	2411f7d: the same (start state and action prepended; agent.py:292-298).
3.3 Policy sampling during imagination	At each imagined step a_i ~ π_θ(· ∣ sg(s_i)), using the current online policy and a stochastic sample. Discrete: Categorical.sample gives an int index. Continuous: mean + std·N(0, 1), an unclipped float sample. The policy input is feat2tensor(sg(carry)) = concat(deter, flatten(stoch)) in the compute dtype. Then deter' = core(deter, stoch, onehot_or_float(a)); logit' = prior(deter'); stoch' ~ OneHot(logit', 0.01). The core's a / sg(max(1, abs(a))) acts as a forward-pass clip to [−1, 1] for continuous actions (2.3). For the loss, the policy is evaluated again on the stop-gradient features with identical parameters, so learning is exactly on-policy. There is no exploration noise beyond the policy's own stochasticity.	Eq. 4 p.5 ("Actor: a_t ~ π_θ(a_t ∣ s_t)"); p.5 imagined trajectories	dreamerv3/agent.py:192 (policyfn = sample(pol(feat2tensor(feat), 1))); dreamerv3/rssm.py:96-101, :137; embodied/jax/outs.py:170-172, :222-224	agree	2411f7d: only out['stoch'] was sg'd inside imgstep, and the start latent and action were sg'd (agent.py:268-293). HEAD sg's the whole carry before the policy. The losses are the same, because features and actions are sg'd wherever they are used. Discrete actions at 2411f7d (re-checked against the code, M6): the 'onehot' actor's sample is a straight-through one-hot (jaxutils.OneHotDist.sample), which imagine consumes directly and the actor loss scores with log_prob of the one-hot; only the env receives an index, nanargmax(act) (agent.py:272, :178). Same values as HEAD's int index one-hot encoded inside imagine.
3.4 Gradient isolation (ac_grads) and stop-gradients	ac_grads = False: imgfeat = concat([sg(first), sg(imgfeat)], axis=1). Actions inside log-probs are sg(act). The cumulative weight is sg(weight). Advantages are sg(adv_normed). Value targets are sg(R) and sg(slowvalue.pred()). So the policy loss updates only policy params and the imagination value loss updates only value params. No actor-critic gradient reaches the RSSM, encoder, reward head or continue head: r̂ and ĉ are computed on stop-gradient features and enter only stop-gradient quantities. The single exception is the replay value loss (3.19), whose input features are not stop-gradiented (repval_grad = True).	Eq. 6 p.6 (sg around the normalised advantage). Detaching imagined features is not stated explicitly, but it is consistent with the paper's "Reinforce for both action types".	dreamerv3/agent.py:196, :411-414, :420-422; dreamerv3/configs.yaml:88, :116	agree	2411f7d: ac_grads 'none', same effect (agent.py:304-308). Re-checked against the code (M6): 2411f7d evaluates the reward and continue heads on the imagined features before the sg (agent.py:301-302, `self.rew(outs).mean()` on `outs`, while `inp = sg(outs)` feeds only the actor and critics), so the statement "r̂ and ĉ are computed on stop-gradient features" holds for HEAD only. Their outputs reach only stop-gradiented quantities (weights, returns, advantages, targets), so no gradient reaches the world model either way; Ajax evaluates them on the sg'd features.
3.5 Reward and continuation used in imagination	r̂i = E[reward head(s_i)], the symmetric two-hot expectation (1.9), a float32 scalar, computed for every i = 0..H. ĉ_i = Binary.prob(1) = sigmoid(con_logit(s_i)), also for every i = 0..H. In the λ-return, step t uses r̂{t+1} and ĉ, predicted from the arrival state s. This matches the world-model convention that reward[t] and is_terminal[t] belong to obs t (2.15). r̂_0 is computed but unused. ĉ_0 IS used, in the weight w_0 (3.6). The continue head's training targets are given in 2.13.	Eq. 1 p.4 (reward and continue predictors on (h_t, z_t)); Eq. 5 p.5	dreamerv3/agent.py:202-206, :174-177; embodied/jax/outs.py:30-31, :189-201; dreamerv3/agent.py:482-490	agree. For the index convention see 3.7.	2411f7d: the start-step reward and continuation came from replay data: startrew = data reward, and startcon = 1 − is_terminal (hard 0/1, without the 0.997 factor). Model predictions were used only for indices 1..H (agent.py:274-302). Since f8817c4, HEAD uses model predictions at every index, so w_t gains an extra factor ĉ_0 ≈ 0.997 for non-terminal starts.
3.6 Discount γ, contdisc, and the cumulative continuation weight	horizon = 333 gives γ = 1 − 1/333 = 0.996997 (the paper's 0.997). contdisc = True folds γ into the continue targets (2.13), so imag_loss uses an explicit disc = 1 if contdisc else 1 − 1/horizon, which is 1 by default. Weight = cumprod(disc·ĉ, axis=1) / disc, with shape (B·K, H + 1). With contdisc this is w_t = Π_{i=0..t} ĉ_i, which includes the start state's ĉ_0, so w_t ≈ 0.997^(t+1) for non-terminal trajectories. Without contdisc it is w_t = γ^t Π_{i=0..t} ĉ_i. The losses use sg(weight[:, :−1]), i.e. w_0..w_{H−1}. A predicted-terminal state s_t (ĉ_t ≈ 0) and every state after it get weight ≈ 0. In repl_loss the discount is always 1 − 1/horizon, because real hard is_terminal flags are used there (3.19).	p.5 (γ = 0.997); Table 4 (1/(1 − γ) = 333). The weighting by cumulative continuation and contdisc are not in the paper.	dreamerv3/agent.py:401-402, :464, :174-177; dreamerv3/configs.yaml:106-107	code_only. Paper Eq. 5/6 sum over t with no continuation weight.	Present at 2411f7d (contdisc True; agent.py:316-317). The difference is that w_0 used the data flag (3.5).
3.7 Imagination λ-return recursion and bootstrap	Call lambda_return(last = zeros, term = 1 − ĉ (soft), rew = r̂, boot = tarval, disc (3.6), λ = 0.95) (Algorithm H). With slowtar = False, tarval = v_t = value.pred()·vscale + voffset, the ONLINE critic's expectation (valnorm 'none' gives offset 0 and scale 1). The config sets slowtar: False, overriding the function default slowtar = True. Recursion: R_H := tarval_H; for t = H−1 down to 0: R_t = r̂{t+1} + disc·ĉ, shape (B·K, H), f32. lambda_return's }·[(1 − λ)·tarval_{t+1} + λ·R_{t+1}]. The output is R_0..R_{H−1`val` argument is unused (only shape-checked); the bootstrap comes entirely from `boot`. The slow critic is not used for targets.	Eq. 5 p.5: R^λ_t = r_t + γc_t((1 − λ)v_t + λR^λ_{t+1}), R^λ_T = v_T; Table 4 (λ = 0.95); p.6 (the EMA regulariser "allows us to compute returns using the current critic network")	dreamerv3/agent.py:397-405, :482-490; dreamerv3/configs.yaml:108	conflict on index convention only. Read literally, Eq. 5 uses r_t, c_t and v_t of the same state; the code uses r̂, ĉ and v of the next state s_{t+1}, the standard DreamerV2 form. Follow the code. Using the online critic for targets agrees with the paper.	The recursion is identical at 2411f7d (agent.py:319-326), also with slowtar False (configs.yaml:146).
3.8 Return normalisation (retnorm)	Normalize(impl 'perc', rate 0.01, limit 1.0, perclo 5, perchi 95, debias False). State: scalar f32 lo = 0 and hi = 0 at init. On each training step, update first, then read: lo ← 0.99·lo + 0.01·percentile(sg(R), 5); hi ← 0.99·hi + 0.01·percentile(sg(R), 95); offset = sg(lo); S = sg(max(1.0, hi − lo)). The percentile is jnp.percentile with linear interpolation over all elements of the flattened (B·K, H) imagination return tensor (15360 values). On multiple devices R is first all-gathered over the data axes. No mask and no weights: steps after a predicted terminal still count. No bias correction, so S stays at its floor of 1 early in training. Only the scale is applied to advantages; the offset is not subtracted (3.10). Updates happen only when training = True. The report pass calls loss(training = False) and does not update.	p.6 Actor learning and Eq. 7 (S = EMA(Per(R, 95) − Per(R, 5), 0.99), denominator max(1, S)); Table 4 (RetNorm scale, limit L = 1, decay 0.99)	embodied/jax/utils.py:16-91 (perc branch :33-35, :51-53, :70-72; _perc :83-88; _update :90-91); dreamerv3/agent.py:407; dreamerv3/configs.yaml:111	agree. Separate EMAs of P5 and P95 equal the EMA of their difference, since both start at 0 with the same rate. Because the update happens before the read, the current batch already contributes to the scale used in the same step.	2411f7d: jaxutils.Moments(impl 'perc'), identical maths, no debias (29eb964:dreamerv3/jaxutils.py:356-386).
3.9 valnorm and advnorm (disabled)	valnorm impl 'none' (rate 0.01, limit 1e-8) and advnorm impl 'none'. Both return (offset 0, scale 1), and their updates are no-ops. So val = value.pred() and critic targets are raw returns. Drop both in Ajax. If enabled, the code path would be val = pred·vscale + voffset, tar_normed = (R − voffset)/vscale and adv_normed = (adv − aoffset)/ascale, with the 'meanstd' impl and debias on. Unused in the recipe.	p.6 (argues against advantage and std-based normalisation)	embodied/jax/utils.py:28-29, :46-47, :63-64; dreamerv3/configs.yaml:112-113; dreamerv3/agent.py:397-399, :409-410, :417-418, :460-462, :468-469	agree	2411f7d: impl 'off', equivalent.
3.10 Advantage	adv_t = (R_t − tarval_t) / S for t = 0..H−1, where tarval_t = tarval[:, :−1] is the online critic's expected value at s_t and S = max(1, hi − lo). The return offset is not subtracted; it would cancel anyway. adv_normed = adv. The policy loss uses sg(adv).	Eq. 6 p.6: sg((R^λ_t − v_ψ(s_t)) / max(1, S))	dreamerv3/agent.py:407-410	agree	Same (2411f7d agent.py:329-332).
3.11 Actor loss (Reinforce for discrete AND continuous)	Per (b, t), t ∈ [0, H): L_pol[b, t] = sg(w_t)·−(logπ_t·sg(adv_t) + η·H_t), where logπ_t = Σ over action keys k of log π_k(sg(a^k_t) ∣ s_t), H_t = Σ_k Entropy[π_k(· ∣ s_t)], and η = actent = 3e-4. This is a pure score-function estimator for both action types. The sampled action is sg, so there is no reparameterised backprop through the dynamics, reward or critic. Vector actions use logp and entropy summed over action dims (Agg sum). Reduction: mean over the H axis, reshape to (B, K), mean, × β_pol = 1.0 (4.1).	Eq. 6 p.6 ("Reinforce estimator for both discrete and continuous actions"); Table 4 (β_pol = 1, η = 3e-4)	dreamerv3/agent.py:411-415, :215, :240; dreamerv3/configs.yaml:86, :108; embodied/jax/outs.py:63-71	agree on the formula. The paper sums over t, the code takes means (4.1). The continuation weight is code_only (3.6).	2411f7d: the same formula (agent.py:333-342), plus a scale_by_actent option (default False) that 1a532a7 removed. 2411f7d keeps the (B·K, H) loss and takes the mean over all its elements (agent.py:342, :414-415) instead of the mean over H then over (B, K): the same value up to float32 rounding.
3.12 Entropy regulariser	η = 3e-4, added inside the weighted bracket. Analytic entropy of the policy at s_t for t = 0..H−1, summed over action keys and dims. Discrete: −Σ p log p from the logits, un-mixed at HEAD (3.14). Continuous: entropy of the unsquashed Normal(tanh(μ), σ) per dim, 0.5·log(2πeσ²), summed over dims. With σ ∈ [0.1, 1] the per-dim entropy lies in [−0.884, 1.419]. No entropy target and no temperature adaptation.	Eq. 6 p.6; p.6 (fixed entropy scale); Table 4	dreamerv3/agent.py:412-414; embodied/jax/outs.py:178-179, :230-234	agree	Same value at 2411f7d (configs.yaml:144). At 2411f7d the discrete entropy was computed on the unimix-mixed distribution.
3.13 Continuous policy distribution ('bounded_normal')	Two output Linears from the last hidden layer, 'mean' and 'stddev', each of action shape, both with outscale 0.01. μ = tanh(mean_raw); σ = (maxstd − minstd)·sigmoid(stddev_raw + 2.0) + minstd = 0.9·sigmoid(x + 2) + 0.1 ∈ (0.1, 1.0). dist = Normal(μ, σ), independent per dim (Agg sum). This is NOT a tanh-squashed Gaussian: samples μ + σ·ε are unbounded, log_prob is the plain Normal logpdf of the unclipped sample, and there is no Jacobian correction. At init (raw outputs ≈ 0): μ ≈ 0 and σ ≈ 0.9·0.8808 + 0.1 = 0.8927. Execution-time bounding: NormalizeAction rescales env bounds to [−1, 1], then ClipAction (outermost) clips the sample to [−1, 1]. Replay stores the unclipped action. The RSSM divides action inputs by max(1, abs(a)). In the code, tanh(mean) and the sigmoid for σ run in the compute dtype (bf16) before the f32 cast in outs.Normal.	Not specified (p.18: 3-layer MLP actor; p.6: Reinforce for continuous actions)	embodied/jax/heads.py:146-155; dreamerv3/configs.yaml:100 (minstd 0.1, maxstd 1.0, outscale 0.01), :103; embodied/jax/outs.py:161-179; dreamerv3/main.py:249-258; embodied/core/wrappers.py:76-86; embodied/core/driver.py:70-77; dreamerv3/rssm.py:137	code_only	2411f7d actor_dist_cont 'normal' had the same maths (tanh mean, 0.9·sigmoid(s + 2) + 0.1, tfd.Normal in Independent) but different precision: raw mean and std outputs were cast to f32 before tanh/sigmoid (2411f7d nets.py:437, :443, :486-488).
3.14 Discrete policy distribution (categorical, unimix)	HEAD: policy_dist_disc = 'categorical'. Head.categorical builds logits = Linear(n_classes) with outscale 0.01 and returns outs.Categorical(logits) without unimix. The configured `policy.unimix: 0.01` is read only by the 'onehot' output and is unused for the policy. Sampling via jax.random.categorical gives an int index; log_prob = log_softmax(logits)[a]; entropy = −Σ softmax·log_softmax. Paper / 2411f7d recipe: probs = 0.99·softmax(logits) + 0.01/n and logits = log(probs). Sampling, log_prob and entropy all use the mixed distribution. Irrelevant for DMC (continuous actions).	p.18 Distributions ("encoder, dynamics predictor, and actor distributions are mixtures of 99% … and 1% uniform"); Table 4 "Actor unimix 1%"; p.20 "1% unimix for all categoricals"	embodied/jax/heads.py:101-110 (categorical, no unimix) vs :112-115 (onehot uses self.unimix); embodied/jax/outs.py:208-234; dreamerv3/configs.yaml:100, :102	conflict: the paper and 2411f7d use 1% actor unimix, HEAD silently uses none.	2411f7d: actor_dist_disc 'onehot' with unimix 0.01 (configs.yaml:128, :130; the nets.py Dist 'onehot' branch, a straight-through one-hot with the argmax index sent to the env). f8817c4 (2024-12-07) changed it to 'categorical', dropping the unimix.
3.15 Policy and critic networks	Input x = concat(deter, reshape(stoch, [..., 32·C])) in the compute dtype, width 10d (12M: 2048 + 512 = 2560). MLP with 3 layers, each Linear(d, bias) → RMSNorm (f32 statistics, eps 1e-4, learned scale, no shift) → SiLU; units d (12M: 256). Init and outscales as in 1.5: policy output 0.01 (mean and std Linears, or logits); value output 0.0 (zero kernel and zero bias, so the logits are exactly 0 and the prediction is exactly 0 at init). The slow critic is an identically shaped MLPHead 'slowval'.	p.18 Networks ("actor and critic … 3-layer MLPs"); p.6 (zero-init critic output); Table 4 (RMSNorm + SiLU); Table 3 p.20 (widths)	dreamerv3/agent.py:51-53, :60-68; embodied/jax/heads.py:16-41, :73-83; embodied/jax/nets.py:144-178, :230-251, :361-399, :565-587; dreamerv3/configs.yaml:100-101, :120-153	agree. The 0.01 policy outscale is code_only.	2411f7d: same init. RMSNorm statistics were in bf16 (1.3).
3.16 Critic output: symexp two-hot	255 symexp bins, raw-space two-hot cross-entropy, symmetric expectation (1.9 and Algorithm D). Targets must be f32. Targets outside ±symexp(20) fall entirely on the edge bin.	Eq. 10-11 p.7; Eq. 12 p.18	embodied/jax/heads.py:132-144; embodied/jax/outs.py:273-330; dreamerv3/configs.yaml:101	agree	Same (29eb964:dreamerv3/jaxutils.py:210-270).
3.17 Critic loss on imagined trajectories (with slow-critic regulariser)	tar = (R − voffset)/vscale = R, padded to H + 1 with a trailing 0. For t ∈ [0, H): L_val[b, t] = sg(w_t)·(CE_twohot(v_ψ(s_t), sg(R_t)) + slowreg·CE_twohot(v_ψ(s_t), sg(E[v_slow(s_t)]))), with slowreg = 1.0. The regulariser target is the slow critic's scalar expectation (slowvalue.pred()), two-hot encoded. It is not a KL or cross-entropy to the slow critic's full distribution. Both terms are evaluated on the same stop-gradient imagined features. State H gets no loss (the pad is sliced off). Reduction: mean over H, then over (B, K), × β_val = 1.0.	Eq. 5 p.5 (max-likelihood loss −Σ ln p_ψ(R^λ_t ∣ s_t)); p.6 (regularise toward an EMA of its own parameters); Table 4 (β_val = 1, critic EMA regulariser 1)	dreamerv3/agent.py:417-422, :215, :240; dreamerv3/configs.yaml:86, :108	agree. The paper does not specify the form of the regulariser. The continuation weighting is code_only.	2411f7d: the same (agent.py:344-350), but the slow critic was computed in f32 (dtype 'float32', agent.py:77-78); HEAD uses the compute dtype.
3.18 Slow critic EMA	SlowModel(rate 0.02, every 1). Init: on first use, the slow params are an exact copy of the online value params (`_initonce`; val(inp) is called before slowval(inp)). Update: once per train step, after the optimizer step: mix = rate if count % every == 0 else 0; θ_slow ← mix·θ + (1 − mix)·θ_slow; count += 1. That is an EMA decay of 0.98 every step. The slow params are not in the optimizer's module list, so they get no gradient. The slow critic is used only as a regularisation target (imagination and replay losses), never for bootstrapping or the advantage baseline (slowtar = False).	p.6 Critic learning (EMA regulariser); Table 4 (critic EMA decay 0.98, regulariser 1)	embodied/jax/utils.py:94-127 (update :113-119, _initonce :121-127); dreamerv3/agent.py:66-68, :74-78, :139-142; dreamerv3/configs.yaml:110	agree	2411f7d SlowUpdater: a separate module with its own random init, computed in f32. mix = clip(1·[updates == 0] + 0.02·[updates % 1 == 0], 0, 1), so a hard copy happens when updater() is first called, which is after the first optimizer step (29eb964:jaxutils.py:737-762; 2411f7d agent.py:77-83, :207-210). HEAD copies at creation and mixes 0.02 at the first update. Equivalent in effect, not by construction: both critics have zero-initialised output layers, so both slow critics predict exactly 0 at step 1, and the warmup gives lr = 0 at step 0, so θ_1 = θ_0.
3.19 Replay value loss (repval)	Enabled by repval_loss = True, scale β_repval = 0.3. Inputs: the replay batch's posterior features repfeat, not stop-gradiented (repval_grad = True), so this loss also trains the encoder and RSSM (2.16). Uses real is_last, is_terminal and reward over the last K = T steps. Bootstrap: boot[b, t] = R^λ_0 of the imagination rollout that started at replay state (b, t), i.e. imag ret[:, 0].reshape(B, K): the imagined λ-return estimate of V(s_t) under the current policy, built from the online critic. Replay λ-return (Algorithm H) with γ = 1 − 1/333 fixed (independent of contdisc) and λ = 0.95: R_{T−1} := boot_{T−1}; for t = T−2..0: live = γ(1 − term_{t+1}), cont = λ(1 − last_{t+1}), R_t = r_{t+1} + live·((1 − cont)·boot_{t+1} + cont·R_{t+1}). At an episode end (last_{t+1}) there is no λ continuation: R_t = r_{t+1} + γ(1 − term_{t+1})·boot_{t+1}. r_{t+1} = `obs['reward'][t+1]`. Loss for t ∈ [0, T−1): L_repval[b, t] = (1 − is_last_t)·(CE_twohot(v_ψ(s_t), sg(R_t)) + 1.0·CE_twohot(v_ψ(s_t), sg(E[v_slow(s_t)]))). The weight is the hard mask ~is_last; terminal steps are also is_last, and the batch annotation marks the step before any mid-window is_first as is_last (5.3). There is no continuation product. The last replay index gets no loss. Reduction: mean over all (B, T−1) = (16, 63) entries including masked zeros, × 0.3. slowtar has no effect here, because tarval goes only into the unused `val` argument.	p.6 Critic learning (replay critic loss β_repval = 0.3, using "the imagination returns R^λ_t at the start states of the imagination rollouts as on-policy value annotations for the replay trajectory to then compute λ-returns over the replay rewards"); Table 4 (0.3)	dreamerv3/agent.py:218-235 (sg skip :220; boot :222), :449-479, :482-490; dreamerv3/configs.yaml:86, :109, :115-116	agree. The ~is_last mask, the explicit γ and the gradient flow into the world model are code details. The gradient flow is implied by the paper's ablation (2.16).	2411f7d: the same (replay_critic_loss True, replay_critic_grad True, replay_critic_bootstrap 'imag', return_lambda_replay 0.95; configs.yaml:104-106, :138; agent.py:352-373).
3.20 Order of operations within one train step	Inside the loss (under nj.grad): (1) world-model forward on replay (posterior repfeat, dyn entries; Algorithm E). (2) Imagine H = 15 steps from all B·T posterior starts with the current policy, and detach the features. (3) Evaluate r̂, ĉ, policy, value and slow value on all H + 1 states. (4) Read the valnorm stats (no-op); compute tarval and the weight; compute R (λ-return); update retnorm with R, then read S; compute the advantage and the policy loss; update valnorm (no-op); compute the value loss. (5) repl_loss: bootstrap from R[:, 0], compute the replay λ-return and the repval loss. (6) Sum the scaled means and take the gradient. After the loss: (7) optimizer update, then the slow-critic EMA update, then latent write-back (5.9). Normaliser updates happen only when training = True.	not stated	dreamerv3/agent.py:137-154 (train), :156-245 (loss), :382-446 (imag_loss), :449-479 (repl_loss)	code_only	Same.

Algorithm H: λ-return (shared by imagination and replay)¶

def lambda_return(last, term, rew, boot, disc, lam):     # all [N, L] f32; the code's unused 'val' arg dropped
    live = (1 - term)[:, 1:] * disc
    cont = (1 - last)[:, 1:] * lam
    interm = rew[:, 1:] + (1 - cont) * live * boot[:, 1:]
    R = [boot[:, -1]]
    for t in reversed(range(L - 1)):
        R.append(interm[:, t] + live[:, t] * cont[:, t] * R[-1])
    return stack(list(reversed(R))[:-1], 1)               # [N, L-1] = R_0 .. R_{L-2}
# imagination: last=0, term=1-c_hat, rew=r_hat, boot=v_online, disc=1 (contdisc), lam=0.95, L=H+1=16
# replay:      last=is_last, term=is_terminal, rew=reward, boot=imag_R0.reshape(B,K), disc=1-1/333, lam=0.95, L=T=64

Algorithm F: imagination actor-critic loss (`imag_loss`)¶

# inputs on (N=B*K, H+1): act (sg), r_hat = rew.pred(), c_hat = con.prob(1), pi, v, v_slow (all on sg feats)
v_on   = v.pred()                                  # online critic expectation
disc   = 1.0                                       # contdisc=True
weight = cumprod(disc * c_hat, axis=1) / disc      # includes c_hat_0
R      = lambda_return(0, 1 - c_hat, r_hat, v_on, disc, 0.95)           # [N, H]
lo, hi = retnorm.update_then_read(sg(R))           # EMA(0.99) of P5 and P95, init 0, no debias
S      = max(1.0, hi - lo)
adv    = (R - v_on[:, :-1]) / S
logpi  = sum_k pi_k.logp(sg(act_k))[:, :-1]
ent    = sum_k pi_k.entropy()[:, :-1]
L_pol  = sg(weight[:, :-1]) * -(logpi * sg(adv) + 3e-4 * ent)
tar    = concat([R, zeros(N, 1)], 1)
L_val  = sg(weight[:, :-1]) * (v.loss(sg(tar)) + 1.0 * v.loss(sg(v_slow.pred())))[:, :-1]
losses.policy = L_pol.mean(1).reshape(B, K); losses.value = L_val.mean(1).reshape(B, K)

Algorithm G: replay critic loss (`repl_loss`)¶

feat = repfeat                                    # NOT sg (repval_grad=True)
boot = R_imag[:, 0].reshape(B, K)                 # imagined lambda-return from each replay state
Rr   = lambda_return(obs.is_last, obs.is_terminal, obs.reward, boot, 1 - 1/333, 0.95)   # [B, T-1]
tar  = concat([Rr, zeros(B, 1)], 1)
L    = f32(~obs.is_last)[:, :-1] * (v(feat).loss(sg(tar)) + 1.0 * v(feat).loss(sg(v_slow(feat).pred())))[:, :-1]
losses.repval = L                                 # [B, T-1]; mean over all entries, x 0.3

4. Loss assembly, gradient routing and optimizer¶

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
4.1 Total loss assembly, per-loss scales and reductions	loss_scales = {rec: 1.0, rew: 1.0, con: 1.0, dyn: 1.0, rep: 0.1, policy: 1.0, value: 1.0, repval: 0.3}. 'rec' is replaced by one entry per decoder observation key, each with scale 1.0. Every term is an un-reduced tensor. World-model terms are asserted to have shape (B, T) = (16, 64): rec_key (summed squared error over the key's feature dims, 2.11), rew (two-hot NLL), con (soft-label Bernoulli NLL), dyn and rep (free-bits KL summed over latents). Imagination terms policy and value are computed on (B·K, H) = (1024, 15), then .mean(axis=1) and reshaped to (B, K) = (16, 64). repval has shape (B, T−1) = (16, 63), masked by (1 − is_last). TOTAL: L = Σ_k scale_k · mean_over_all_elements(loss_k). It is a mean, not divided by the sum of weights, so weighted-out or masked entries still count in the denominator. World-model terms are not masked at episode boundaries; the RSSM handles resets through is_first. The world-model and actor-critic losses form one f32 scalar, differentiated w.r.t. all modules [dyn, enc, dec, rew, con, pol, val] and applied in one optimizer step.	Eq. 2 p.4 (L = E[Σ_{t=1..T}(β_pred L_pred + β_dyn L_dyn + β_rep L_rep)]); Eq. 5/6 (sums over t); Table 4 (β_pred = 1, β_dyn = 1, β_rep = 0.1, β_val = 1, β_repval = 0.3, β_pol = 1)	dreamerv3/agent.py:80-83 (scales, rec expanded per key), :184-186 (shape assert), :215 (imagination reshape), :219-235 (repval), :237-240 (sum of scaled means), :74-78 (single optimizer); dreamerv3/configs.yaml:86	conflict in reduction form, while the scale values agree. The paper sums over t; the code takes means over B × T and over H. A sum would make the world-model terms 64× and the actor-critic terms 15× larger, which changes their relative weighting. Global rescaling is otherwise absorbed by LaProp (eps 1e-20) except through AGC thresholds. Reproduce the code's means.	2411f7d: scales {dec_cnn 1, dec_mlp 1 (per vector key), reward 1, cont 1, dyn 1, rep 0.1, actor 1, critic 1, replay_critic 0.3} with the same mean reduction (agent.py:414-415), written Σ_k mean(scale_k · loss_k): each element is scaled before the mean (the same value up to float32 rounding). v1: separate model, actor and critic optimizers.
4.2 Gradient routing inside the single optimizer (summary)	See the routing matrix below; details in 2.16 and 3.4. reward_grad = True; the continue and decoder inputs are not sg; repval_grad = True (the replay critic trains enc and dyn); ac_grads = False (imagination starts and imagined feats are sg, so policy and value losses update only pol and val). Inside imagine, the policy sees sg(carry). After the optimizer step: slowval ← 0.02·val + 0.98·slowval (3.18). The retnorm percentile EMA is updated inside the loss forward pass (training = True), once per step.	p.4 (end-to-end world model); Eq. 6 (sg advantage); Ablations p.10, Fig. 6b, Fig. 18 (reward and value gradients shape the representation); Table 4 (critic EMA decay 0.98, regulariser 1)	dreamerv3/agent.py:139-142, :172, :196, :219-220; dreamerv3/configs.yaml:88, :110, :114-116; embodied/jax/utils.py:94-127	agree (implicit for repval, see 2.16)	2411f7d: replay_critic_grad True, reward_grad True, ac_grads 'none', slow fraction 0.02 every 1. Same.
4.3 Optimizer: one LaProp chain over ALL trainable modules	One optimizer instance over modules [dyn, enc, dec, rew, con, pol, val]: world model, actor and critic in one param tree, one scalar loss, one update per train step. Not optimised: slowval (EMA copy of val) and the retnorm, valnorm and advnorm statistics. Per step: (L, grads) = value_and_grad of L w.r.t. the f32 params (loss asserted f32 with shape ()); on multiple devices grads = pmean over the data axes. updates = chain(grads), with the chain in this exact order: (1) clip_by_agc(0.3, pmin = 1e-3) per tensor (4.4); (2) scale_by_rms(beta2 = 0.999, eps = 1e-20), bias-corrected (4.5); (3) scale_by_momentum(beta1 = 0.9, nesterov = False), bias-corrected (4.6); (4) add_decayed_weights only if wd > 0 (wd = 0); (5) optax.scale_by_learning_rate(schedule), i.e. × −lr_k (4.7). params ← params + updates. Optimizer state is f32. The config key opt.momentum = True is accepted by _make_opt but ignored: momentum is always appended. See Algorithm I.	p.18 Optimizer; Table 4 (LR 4e-5, AGC(0.3), LaProp(ε = 1e-20)); p.20 "Previous Dreamer generations" (AGC, LaProp = RMSProp followed by momentum)	dreamerv3/agent.py:74-78, :342-379 (_make_opt; chain :358-360); embodied/jax/opt.py:16-81 (f32 scalar assert :37-38; pmean :52-54; update :60-63); dreamerv3/configs.yaml:87	agree, apart from beta2 (4.5)	2411f7d: the same single optimizer (separate_lrs False) in dreamerv3/jaxutils.py, with an optional globclip (0.0 = off) before AGC. Identical maths.
4.4 Adaptive gradient clipping (AGC), per TENSOR	For every parameter leaf p with gradient g (each kernel, bias and norm scale is its own tensor): u = ‖vec(g)‖₂, n = ‖vec(p)‖₂ (current, pre-update params), upper = 0.3·max(1e-3, n), g ← g / max(1, u/upper). Applied to the raw gradient of the summed loss, before RMS normalisation. This is not NFNet unit-wise AGC: there are no per-row norms, and the norm runs over the whole flattened tensor, including 3-D BlockLinear kernels. Zero-initialised tensors (all biases, the zero output kernels of the reward and value heads) have n = 0, so their gradient norm is capped at 3e-4 until they grow. This cap does not slow those tensors: a constant per-tensor factor cancels in the following g / sqrt(EMA(g²)) with eps = 1e-20, and at the first RMS step u = sign(g) whatever the clip. As ‖p‖ grows the cap loosens, which makes normalised updates temporarily larger while the EMA catches up. pmin = 1e-3 is a hard-coded default at HEAD.	p.18 ("clips per-tensor gradients if they exceed 30% of the L2 norm of the weight matrix they correspond to, with its default ε = 10⁻³")	embodied/jax/opt.py:109-123; dreamerv3/agent.py:358; dreamerv3/configs.yaml:87 (agc 0.3)	agree	2411f7d jaxutils.scale_by_agc: identical formula, pmin configurable (1e-3).
4.5 LaProp normaliser (RMS step), beta2 conflict	State: step (int32, init 0) and nu (f32 zeros). step ← step + 1; nu ← β₂·nu + (1 − β₂)·g² (g is the AGC-clipped gradient); nu_hat = nu / (1 − β₂^step); u = g / (sqrt(nu_hat) + eps), with eps added after the sqrt. Code: β₂ = 0.999, eps = 1e-20. With such a small eps the update is invariant to a global rescaling of the loss.	p.18: "LaProp optimizer with ε = 10⁻²⁰ and its default parameters β₁ = 0.9 and β₂ = 0.99" (verified on PDF p.18)	embodied/jax/opt.py:126-143; dreamerv3/agent.py:348, :359; dreamerv3/configs.yaml:87 (beta2 0.999)	conflict: the paper says 0.99, the code 0.999	2411f7d: {scaler rms, beta2 0.999, eps 1e-20}, the same as HEAD. So the published curves were produced with 0.999. float32 caveat (observed running 29eb964 in M6: one step of 195 learning rates, from g ≈ 2e-18): the term that underflows is the increment (1 − β₂)·g². Below the smallest normal float32 (1.18e-38), i.e. for
4.6 LaProp momentum step	Its own step counter (int32) and mu (f32 zeros). step ← step + 1; mu ← β₁·mu + (1 − β₁)·u, where u is the RMS-normalised update from 4.5; output mu_hat = mu / (1 − β₁^step). β₁ = 0.9, nesterov = False. Momentum acts on the normalised update (LaProp), not on the raw gradient (Adam).	p.18 ("LaProp normalizes gradients by RMSProp and then smoothes them by momentum"), β₁ = 0.9	embodied/jax/opt.py:146-164; dreamerv3/agent.py:360	agree	Identical in 1a532a7 jaxutils.scale_by_momentum.
4.7 Learning rate and warmup schedule	lr = 4e-5. schedule 'const', anneal 0, warmup 1000 updates: ramp = optax.linear_schedule(0, lr, 1000); sched = join_schedules([ramp, constant(lr)], [1000]); applied through optax.scale_by_learning_rate (sign flipped). With the optax count k = 0, 1, 2, … (read before increment): lr_k = 4e-5 · k/1000 for k < 1000, then 4e-5. The first update uses lr = 0, so it is a no-op on the params, but the moments still accumulate. No decay afterwards. The warmup covers the actor and critic too, since they share the optimizer. The unused 'linear' and 'cosine' options decay to 0.1·lr over (anneal − warmup) steps.	Table 4 (LR 4e-5); Table 4 caption: "We do not use any hyperparameter annealing"	dreamerv3/agent.py:366-378; dreamerv3/configs.yaml:87 (schedule const, warmup 1000, anneal 0)	code_only for the warmup (the paper describes none); the constant LR agrees	2411f7d: updates *= clip(step/1000, 0, 1) with step read before increment, giving numerically identical lr_k.
4.8 Disabled regularisation and clipping	Weight decay wd = 0.0. If enabled it would be optax.add_decayed_weights(wd) on keys matching r'/kernel$', after momentum and before the LR scaling. No global-norm gradient clipping, no dropout, no prioritised replay, no annealing.	Table 4 caption p.21 ("We do not use any hyperparameter annealing, prioritized replay, weight decay, or dropout")	dreamerv3/agent.py:351-365; dreamerv3/configs.yaml:87	agree	2411f7d had globclip 0.0 (off) before AGC.
4.9 fp16 dynamic loss scaling (inactive)	Only when the compute dtype is float16: loss = grad_scale (init 1e4); grads = 1/grad_scale; optax.apply_if_finite(max_consecutive_errors = 1000); grad_scale ×2 after 1000 consecutive finite steps and /2 on a non-finite step, clipped to [1e-4, 1e5]; the optimizer step counter increments only on finite steps. None of this runs with the default bfloat16. Skip it.	not in paper	embodied/jax/opt.py:25-29, :39-46, :56-70, :83-92	code_only	Identical in 1a532a7 jaxutils.Optimizer.

Gradient routing matrix (which loss updates which module)¶

"RSSM" is the code module dyn: core, prior MLP and posterior obs layers. The straight-through mark (ST) means the gradient reaches the posterior logits through z_t.

Loss (scale)	enc	RSSM core	RSSM prior MLP	RSSM posterior layers	dec	rew	con	pol	val	slowval
rec per key (1.0)	yes	yes	no	yes (ST)	yes	no	no	no	no	no
rew (1.0)	yes	yes	no	yes (ST)	no	yes	no	no	no	no
con (1.0)	yes	yes	no	yes (ST)	no	no	yes	no	no	no
dyn (1.0)	yes, via h_t and earlier steps	yes	yes	via earlier z through h_t	no	no	no	no	no	no
rep (0.1)	yes	yes	no	yes	no	no	no	no	no	no
policy (1.0)	no	no	no	no	no	no	no	yes	no	no
value (1.0)	no	no	no	no	no	no	no	no	yes	no
repval (0.3)	yes	yes	no	yes (ST)	no	no	no	no	yes	no
(EMA after step)										0.02·val + 0.98·slowval

Algorithm I: one optimizer update (shared by all modules)¶

def opt_update(params, grads, state):                 # all f32
    k = state.count                                   # 0, 1, 2, ...
    g = {name: gi / max(1.0, norm(gi) / (0.3 * max(1e-3, norm(params[name]))))
         for name, gi in grads.items()}               # AGC per tensor
    state.t_rms += 1
    state.nu = 0.999 * state.nu + 0.001 * g**2
    u = g / (sqrt(state.nu / (1 - 0.999**state.t_rms)) + 1e-20)
    state.t_mom += 1
    state.mu = 0.9 * state.mu + 0.1 * u
    m = state.mu / (1 - 0.9**state.t_mom)
    lr = 4e-5 * min(k / 1000, 1.0)                    # first update: lr = 0
    state.count += 1
    return params - lr * m, state
# then: slowval = 0.02 * val + 0.98 * slowval

5. Replay buffer and replay context¶

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
5.1 Storage layout and capacity	There is one append-only stream per env worker. Episodes are concatenated in the stream, and boundaries exist only as the is_first / is_last flags. Each stored step t contains: obs_t (vector keys, reward_t, is_first_t, is_last_t, is_terminal_t); the action a_t sampled after observing obs_t, zeroed by the driver if is_last_t; the policy outputs 'dyn/deter' = h_t (f32[D]) and 'dyn/stoch' = z_t (f32[32, C] one-hot) (5.7); and 'stepid' (uint8[20] = a 16-byte chunk UUID plus a 4-byte big-endian index). Storage is in chunks of 1024 steps linked by `succ`. An item is a start position. When step n of a worker is added and the worker's stream holds at least `length` steps, the item starting at n − length + 1 is inserted. Every step becomes the start of exactly one item (stride-1 overlapping windows). length = consec_train·batch_length + replay_context = 1·64 + 1 = 65. capacity = 5e6 items (not episodes). When full, the oldest item is evicted FIFO, and chunk memory is freed by refcount. A capacity check asserts batch_size·length ≤ capacity.	p.18 "Experience replay"; Table 4 (replay capacity 5e6); p.20 ("Larger capacity, online queue, storing and updating latent states")	embodied/core/replay.py:77-118 (add), :171-191 (insert/remove FIFO), :193-214 (getseq across chunks); embodied/core/chunk.py:41-62; dreamerv3/main.py:183-209; dreamerv3/configs.yaml:39-46	agree	2411f7d: replay length = batch_length = 65 (context included), same item semantics (1a532a7 embodied/replay/replay.py).
5.2 Sliding windows cross episode and chunk boundaries	Items are the 65-step windows ending at each insert in a worker's continuous stream. They are not aligned to episodes: windows freely cross episode boundaries and chunk boundaries. A mid-sequence is_first is the normal case, and the RSSM resets the carry there (2.2). Do not restrict sampling to within-episode windows.	p.18 (not explicit)	embodied/core/replay.py:108-118	code_only (the paper is silent)	Same at 2411f7d.
5.3 Training batch sampling and annotation	B = 16 rows, sampled independently and with replacement. For each row in 'train' mode: if the online queue is non-empty, pop its oldest item (5.4); otherwise draw uniformly over all current items (Uniform selector: rng.integers(0, len(keys))). Gather 65 contiguous steps into arrays [16, 65, ...]. Annotate: is_first[:, 0] = True; is_last = is_last OR roll(is_first, −1), with the last column of the rolled array set to False. This marks the step before any mid-window is_first as last, which also covers abandoned episodes. The Consec stream (consec = 1, prefix = replay_context = 1, strict) emits the whole [16, 65] window with consec = 0 for every step. A Prefetch thread (amount = 1) prepares the next batch (6.3). With replay_context = 1, is_first[:, 0] = True has no effect on training, because index 0 is only context and its obs are dropped. It matters when replay_context = 0 (5.11).	p.18: "each minibatch is formed first from non-overlapping online trajectories and then filled up with uniformly sampled trajectories from the replay buffer"	embodied/core/replay.py:121-127 (sample), :151-169 (_sample), :255-275 (assemble), :277-292 (annotate); embodied/core/selectors.py:29-57; embodied/core/streams.py:89-140 (Consec), :32-86 (Prefetch); dreamerv3/main.py:261-272; embodied/jax/agent.py:326-337	agree	1a532a7 replay.dataset: same sampling and annotation; there the report dataset also popped the online queue (HEAD: train mode only).
5.4 Online queue	replay.online = True. Each worker has a counter lengths[w], the number of steps it has previously added (incremented after each add). When step n (0-indexed) inserts an item (n ≥ 64) and lengths[w] % 65 == 0, i.e. n = 65, 130, …, the newly inserted item (start n − 64 = 1, 66, 131, …) is appended to a FIFO deque. This yields non-overlapping consecutive 65-step windows of the freshest data. The queue is popped only by 'train' sampling (report never pops) and is unbounded. With 16 envs it gains 16/65 ≈ 0.246 items per vector step, against 16·(train_ratio/1024)·16 rows sampled per vector step. At RR 512 that is 128 rows per vector step, so about 0.2% of rows are online (about 3% at RR 32).	p.18 ("uniform replay buffer with an online queue")	embodied/core/replay.py:39-42, :108-118, :158-160	agree	Same push rule at 2411f7d and 1a532a7 (embodied/replay/replay.py:140-144).
5.5 Non-uniform selectors (present but disabled)	fracs {uniform 1.0, priority 0.0, recency 0.0}. A Mixture is built only if fracs.uniform < 1. Prioritized(exponent 0.8, initial inf, zero_on_sample True, maxfrac 0.5, branching 16) and Recency(1/(age + 1)^1.0, branching 16) exist in the code. HEAD's agent never writes priorities (agent.py:151-152 is commented out). Implement uniform only.	p.18 ("we opt for uniform replay"); Table 4 caption (no prioritised replay)	dreamerv3/configs.yaml:42-45; dreamerv3/main.py:197-207; embodied/core/selectors.py:60-229; dreamerv3/agent.py:151-152	agree	2411f7d wrote priorities when fracs.priority > 0 (priosignal 'model' = rep + dyn + reconstruction loss per step). Off by default.
5.6 Eviction handling (sampling retry, partial write-back)	If a popped online item or a uniformly drawn item points to an evicted chunk (KeyError in _getseq), `_sample` silently retries: the next queue item, or a fresh uniform draw. Latent write-back for rows whose chunk was evicted is skipped. If a successor chunk is missing, the write can stop after a partial update.	none	embodied/core/replay.py:156-169, :141-149, :216-235	code_only	Same.
5.7 Replay context: storing latents at collection	replay_context = 1. At every policy step, the RSSM posterior entry for each env, entry_t = {deter: h_t, stoch: z_t}, is returned as outs {'dyn/deter', 'dyn/stoch'} and stored with transition t. For vector obs the enc and dec entries are empty. entry_t is the state after filtering obs_t: h_t = core(h_{t−1}, z_{t−1}, a_{t−1}) (with is_first masking), and z_t is a sample (not the mode) from OneHot(logits, 0.01). Agent.ext_space adds the replay keys consec (int32), stepid (uint8[20]), dyn/deter (f32[D]) and dyn/stoch (f32[32, C]). A memory-efficient lossless equivalent is z as int32 class indices [32] and h in bf16.	p.18 ("We store latent states into the replay buffer during data collection to initialize the world model on replayed trajectories")	dreamerv3/agent.py:89-99 (ext_space), :115-135; dreamerv3/rssm.py:39-43 (entry_space), :75-92; embodied/core/driver.py:70-80	agree	1a532a7 and 2411f7d stored stoch as int32 argmax indices and one-hot encoded them at train time (algorithmically identical).
5.8 Replay context: RSSM initialisation at the start of a training sequence	Data per row has T + K = 64 + 1 = 65 steps. `_apply_replay_context` with K = 1 takes the 'replay' branch whenever `data['consec'][:, 0]` == 0, which is always the case with consec_train = 1. carry = stored entries at index K − 1 = 0 (dyn.truncate takes `entries[:, :K][:, −1]`): (h_0, z_0) from the buffer, a constant with no gradient. obs = data[:, 1:65]; prevact = `data[action][:, 0:64]` (actions at indices 0..63 = a_{t−1} for each obs_t, so the first trained step consumes a_0); stepid = stepid[:, 1:]. The training scan then computes h_1 = core(h_0, z_0, a_0), unless is_first_1 holds (then the state and action are zeroed). Step 0 contributes no loss. Stored latents come from older params and are used as-is. See Algorithm J.	p.18 ("…to initialize the world model on replayed trajectories")	dreamerv3/agent.py:137-140, :312-340; dreamerv3/rssm.py:51-54 (truncate); dreamerv3/main.py:187; embodied/jax/agent.py:420, :447-457	agree	29eb964 (2024-05-15, after 2411f7d) fixed a bug: the first prevact came from the train carry, i.e. the last action of the previous, unrelated batch, instead of data[:, K − 1]. HEAD has the fix (rep_prevact = data[:, K−1:−1]). So the paper-era code 2411f7d still had the bug.
5.9 Latent write-back after each update	train returns outs['replay'] = {stepid [B, T, 20] (indices 1..64), 'dyn/deter' [B, T, D], 'dyn/stoch' [B, T, 32, C]}. These are the posterior entries of the training forward pass: pre-update params, sampled posteriors, after replay-context initialisation. replay.update: for each row i, decode (chunk UUID, index) from stepid[i, 0] and overwrite the latents of the 64 consecutive steps, crossing chunk boundaries via succ. Rows whose chunk was evicted are skipped (5.6). The context step (index 0) is not overwritten by this row. Report passes never write back. Each step is refreshed roughly train_ratio times over its life. Asynchrony in the reference: the jax Agent wrapper returns the outs of the previous train call (pending_outs is fetched asynchronously). So the first update writes nothing back, update k's latents reach the buffer only after update k + 1 has run, and the final update's latents are never written. A synchronous Ajax implementation that writes back immediately is a faithful simplification (Open questions).	p.18 ("and write the fresh latent states of the training rollout back into the buffer")	dreamerv3/agent.py:143-150; embodied/run/train.py:76-79; embodied/core/replay.py:130-149, :216-235; embodied/jax/agent.py:286-289	agree (the delay is code_only infrastructure)	Same at 2411f7d, including the delayed outs (2411f7d:dreamerv3/jaxagent.py:188-191).
5.10 Train carry threading and consec_train	carry_train = (enc {}, dyn {deter, stoch} zeros, dec {}, prevact zeros) from init_train(B) is threaded through successive train calls. train returns (the final enc/dyn/dec carry of the trained window, computed with pre-update params; `data[act][:, −1]`). It is consumed only when consec_train > 1. In that case one sampled sequence of length consec·T + K is split into chunks [kT, kT + T + K) with consec = k. For k ≥ 1 the 'normal' branch uses the threaded carry, whose state corresponds to index K − 1 of chunk k, with obs/prevact = data[:, K:] / `data[act][:, K−1:−1]`. Gradients never cross chunk boundaries. Default consec_train = 1, so the threaded carry is always overridden by the replay context, and this machinery can be omitted.	not described	embodied/run/train.py:67, :76; dreamerv3/agent.py:109-110, :137-154, :312-340; embodied/core/streams.py:89-140; dreamerv3/configs.yaml:13-15	code_only	2411f7d: the carry was passed similarly; long-sequence splitting was replay.dataset(length) in 1a532a7.
5.11 replay_context = 0 behaviour (resolved in the merge)	With replay_context = 0, `_apply_replay_context` returns early. The carry passed in is the previous train call's carry and prevact = concat(carry prevact, `data[act][:, :−1]`). However, Replay.sample always sets is_first[:, 0] = True, and `_observe` masks deter, stoch and prevact to zero wherever is_first holds. So the sequence starts from h_0 = core(0, 0, 0): a zero-state (DreamerV2-style) start, not a reuse of unrelated state. Use this as the fallback if Ajax's replay cannot store latents: run with replay_context = 0 and force is_first at index 0.	p.18 (the paper describes only the stored-latent version)	embodied/core/replay.py:126, :284-286; dreamerv3/rssm.py:75-80; dreamerv3/agent.py:316-320	code_only	This corrects the world-model extractor's note, which claimed the previous batch's carry is reused.

Algorithm J: one training step with the replay context (consec_train = 1)¶

# data: dict of [B, 65, ...] incl. obs keys, action, reward, is_first/last/terminal, dyn/deter, dyn/stoch, stepid
carry0  = dict(deter=bf16(data['dyn/deter'][:, 0]), stoch=bf16(data['dyn/stoch'][:, 0]))   # constants, no grad
obs     = {k: data[k][:, 1:] for k in obs_keys}            # [B, 64]
prevact = {k: data[k][:, 0:64] for k in act_keys}          # a_{t-1} for each obs_t
(loss, aux), grads = value_and_grad(total_loss)(params, carry0, obs, prevact)   # Algorithms E, F, G
params, opt_state = opt_update(params, grads, opt_state)   # Algorithm I
slowval = 0.02 * params.val + 0.98 * slowval
replay.update(stepid=data['stepid'][:, 1:],                # write back fresh posteriors for indices 1..64
              deter=aux.entries.deter, stoch=aux.entries.stoch)   # computed with the PRE-update params

6. Training loop, scheduling and infrastructure¶

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
6.1 Train-ratio semantics and update scheduling	`step` counts agent transitions summed over all env workers: each env contributes 1 per vector step, and action repeat is not counted. batch_steps = B·T = 16·64 = 1024 (context step excluded). should_train = Ratio(train_ratio / 1024), where Ratio(r)(step) is: if prev is None, set prev = step and return 1; otherwise repeats = floor((step − prev)·r), prev += repeats/r, return repeats. It is evaluated after every single transition. The per-transition callback order is step += 1, then replay.add, then trainfn, so updates interleave with the per-env adds inside a vector step. Gradient updates per agent step = train_ratio/1024; per env frame = train_ratio/(1024·repeat). For B = 16 and T = 64: RR 32 gives 1/32 updates per agent step (Atari, repeat 4: 1 update per 128 frames, about 1.56M updates over 200M frames); RR 64 gives 1/16; RR 128 gives 1/8; RR 256 gives 1/4; RR 512 gives 1/2 (DMC proprio: 250K agent steps × 0.5 = about 125K updates); RR 1024 gives 1. The logged replay stat 'replay_ratio' = 65·samples/inserts = (65/64)·train_ratio. The separate 'parallel' script enforces the ratio with SamplesPerInsert instead (train_ratio/batch_length samples per inserted step, tolerance 4·B, minsize B·65).	p.18: replay ratio = "fraction of time steps trained on per time step collected from the environment, without action repeat. Dividing the replay ratio by the time steps in a minibatch and by action repeat yields the ratio of gradient steps to env steps" (Atari example: RR 32, repeat 4, 16 × 64 gives 1 step per 128 env steps, 1.5M over 200M)	embodied/run/train.py:24-25, :59-62, :70-81; Ratio: 2411f7d:embodied/core/when.py:26-42 (HEAD imports elements.when.Ratio from the external `elements` package, not vendored); embodied/core/replay.py:70; embodied/run/parallel.py:242-245	agree	2411f7d: batch_steps = B·(batch_length − replay_context) = 16·(65 − 1) = 1024, identical.
6.2 Learning starts, prefill and exploration	No random-action prefill and no exploration schedule. The randomly initialised actor acts from step 0 (continuous: mean ≈ 0, std ≈ 0.893). The gate: trainfn returns early until len(replay) (items) ≥ B·T = 1024. A worker yields items only after 65 steps, so with N lock-stepped envs there are N·(s − 64) items after s steps per env. With N = 16 the first update happens at s = 128, i.e. 2048 total transitions. That first eligible Ratio call returns 1 (one update); afterwards accrual is normal, with no catch-up for the prefill period. Prefetch subtlety (reference only): the train stream's Prefetch worker starts when the stream is created at setup and calls replay.sample immediately. replay.sample waits only until len(sampler) ≥ 1, not for the 1024-item gate. So the first training batch is drawn as soon as about 1 to 16 items exist (about 65 steps per env), uniformly with replacement from the earliest items, with heavy duplication. The online queue is usually still empty at that point, because the first per-worker push is at n = 65; this part is race-dependent. When training starts (2048 transitions), each worker has pushed exactly one queue item (n = 65; n = 130 is not reached), so the second batch, sampled right after the first is dequeued, is exactly the 16 online windows starting at index 1. In general batch k + 1 is sampled when batch k is dequeued, before update k runs.	not specified	embodied/run/train.py:64 (stream built at setup), :70-73 (gate), :93-95; embodied/jax/agent.py:337 (Prefetch); embodied/core/streams.py:44-48, :50-53, :70-77; embodied/core/replay.py:114, :121-124; 2411f7d:embodied/core/when.py:33-42	code_only	2411f7d: training started when len(replay) ≥ batch_size (16 items) and step ≥ train_fill (0), with expl_until 0 (no explore mode), so the first update came after about 65 steps per env.
6.3 Asynchrony and staleness in the reference loop (engineering, not algorithmic)	(a) Acting params: only keys matching '^(enc\|dyn\|dec\|pol)/' are copied to the policy. A train call sets pending_sync to its pre-update non-donated params (`allo`), but only if no sync is already pending; otherwise `allo` is deleted. The swap happens only after the next policy call has already run on the old params. Let θ_v be the params after all updates of driver iteration v. Policy call π_{v+1} then acts with θ_{v−2} while the trainer is at θ_v, so the acting params lag by about two driver iterations of updates (about 16 updates at RR 512 with 16 envs). (b) Batch prefetch (amount 1) samples the next batch before the current update and its latent write-back (6.2). (c) Delayed outs: both metrics and outs['replay'] are returned one train call late (5.9). None of this changes the algorithm in expectation. A synchronous Ajax loop (act with the current params, sample after the write-back, write back immediately) is a faithful simplification, but it will not reproduce the reference bit-for-bit.	not in paper	embodied/jax/agent.py:143-147, :229-247 (policy then swap), :264-289 (allo / pending_sync / pending_outs); embodied/jax/transform.py:78-89; embodied/core/streams.py:32-86; embodied/core/driver.py:70-80; dreamerv3/agent.py:85-87	code_only	2411f7d: policy keys '/(enc\|dyn\|actor)/' with a similar delayed sync, and the same delayed outs (2411f7d jaxagent.py:188-191).
6.4 Report / diagnostics stream (not part of learning)	report_length 32. The report stream samples from the train replay, whose items are 65 steps long, in 'report' mode: it never pops the online queue and is not counted in the sample stats. Annotation (is_first[:, 0] = True, is_last OR next is_first) is applied to the full 65-step window. Consec(length 32, prefix 1, strict False) then slices [:, 0:33], so report windows are the first 33 steps of a 65-step train item. consec_report 1, report_batches 1, run every report_every = 300 s. It runs the loss with training = False: no normaliser updates, no param update, no latent write-back. For image keys it also makes open-loop prediction videos. Optional for Ajax.	not in paper	embodied/run/train.py:65, :68, :99-104; dreamerv3/main.py:261-270 (strict only for train); embodied/core/streams.py:126-133; dreamerv3/agent.py:247-310; dreamerv3/configs.yaml:12, :14, :53, :58	code_only	2411f7d: batch_length_eval 33, eval_every 180 s.
6.5 RNG seeding	Param init seed = [config.seed, 0]. Policy seed per call = default_rng([config.seed, n_actions]).integers(0, 2³² − 1, (2,), uint32). Train seed per batch = the same construction with the counter n_batches. That counter is shared by the train and report Prefetch streams (both go through agent.stream), so train-batch seeds depend on how many report batches were prefetched. The replay Uniform selector rng = np.random.default_rng(0), not derived from config.seed (Replay(seed = 0) default; main.py never passes a seed); the Mixture rng also uses seed 0. XLA deterministic ops are enabled. Ajax: seed the replay sampler from config.seed and document the deviation (Open questions).	not in paper	embodied/jax/agent.py:332-335 (shared n_batches), :405-408, :418-419; embodied/run/train.py:64-65; embodied/core/replay.py:16-25; embodied/core/selectors.py:31-34, :202-212; embodied/jax/internal.py:45-47	code_only	—

Algorithm K: reference training loop (synchronous simplification recommended for Ajax)¶

step = 0
ratio = Ratio(train_ratio / (B * T))                  # 512 / 1024 = 0.5 for the DMC-proprio paper setting
obs = envs.reset()                                    # N = 16 envs; is_first=True, reward=0
carry = init_policy(N)                                # zeros; prevact zeros
while step < total_agent_steps:                       # 250K agent steps = 500K env steps at repeat 2
    carry, act, entries = policy(params, carry, obs)  # Algorithm B on the (N,) batch, then sample pi
    act_env = clip(act, -1, 1)                        # ClipAction then NormalizeAction
    act_store = where(obs.is_last, 0, act)            # driver zeroes actions at is_last
    for w in range(N):                                # per transition
        step += 1
        replay.add(w, {**obs[w], 'action': act_store[w], 'dyn/deter': entries.deter[w],
                       'dyn/stoch': entries.stoch[w]})
        if len(replay) >= B * T:                      # items, 1024
            for _ in range(ratio(step)):
                train_step(replay.sample(B))          # Algorithm J; online-queue pops first
    obs = envs.step(act_env, reset=obs.is_last)       # is_last envs are reset; next obs is_first

7. Acting, environment protocol, benchmark configurations and sanity targets¶

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
7.1 Act-time computation (collection policy)	Batched over N envs per vector step. carry = (enc {}, dyn {deter [N, D] bf16, stoch [N, 32, C] bf16}, dec {}, prevact). reset = obs['is_first']. tokens = enc(obs) (symlog vector obs, then the 3-layer MLP). RSSM._observe (Algorithm B): deter, stoch and prevact are masked to 0 where reset; prevact is one-hot encoded (discrete) and masked again; h_t = core(...); posterior logits from concat(h_t, tokens); z_t sampled from OneHot(logits, 0.01). feat = concat(h_t, flatten(z_t)). a_t ~ π(· ∣ feat), SAMPLED. New carry = (..., {deter: h_t, stoch: z_t}, ..., a_t as the raw sample). The decoder is not run at act time. Outputs: actions to the env, entries to the replay (5.7). The carried prevact is the raw sample even at is_last steps (the driver zeroes only the stored and sent action). This is harmless, because the next obs is is_first and masking zeroes it.	Eq. 1 p.4; p.5 ("For environment interaction, we select actions by sampling from the actor network without lookahead planning")	dreamerv3/agent.py:101-107, :115-135; dreamerv3/rssm.py:61-92, :135-159; embodied/jax/agent.py:219-261; embodied/core/driver.py:56-83	agree	1a532a7 policy: the same filter and sample.
7.2 is_first resets and episode protocol (driver)	The initial acts are zeros with reset = True for all envs. env.step({..., reset: True}) returns the first obs (is_first = True, reward = 0, is_terminal = False). After each policy call, actions for envs with is_last are multiplied by 0, so a zero action is stored in the replay at the is_last step. reset = True is sent at the next step, whose obs is is_first of a new episode in the same worker stream. RSSM masking is described in 2.2. is_terminal = (dm_env discount == 0) and is False on first steps. DMC time-limit ends give is_last = True and is_terminal = False.	not detailed	embodied/core/driver.py:35-40, :70-80; dreamerv3/rssm.py:75-80; embodied/jax/nets.py:67-77; embodied/envs/from_dm.py:51-78	code_only	1a532a7 nets.RSSM.observe: jaxutils.reset((deter, stoch, action), reset) before _gru. Same.
7.3 Eval mode	Agent.policy ignores its mode argument. 'train', 'eval' and 'report' all sample the actor (no mean or argmax action) and sample posteriors. The default 'train' script logs episode/score from the training envs (stochastic policy, 16 envs). train_eval adds eval envs (eval_envs 4, eval_eps 1 per eval) using the same sampling policy and a separate eval replay of capacity size/10.	not specified	dreamerv3/agent.py:115, :126; embodied/run/train.py:93; embodied/run/train_eval.py:126-134; embodied/run/eval_only.py:64; dreamerv3/main.py:186; dreamerv3/configs.yaml:56-57	code_only	1a532a7 Agent.policy also ignored mode (always sample).
7.4 Action handling at the env boundary	Continuous ('bounded_normal', 3.13): sample a = μ + σ·ε in f32, not squashed. Env side: wrap_env applies NormalizeAction (maps [−1, 1] to [low, high]) and then ClipAction (outermost, clips to [−1, 1]) to every continuous key. The replay stores the UNCLIPPED sample. RSSM input: a / sg(max(1, abs(a))). Discrete ('categorical' at HEAD, 3.14): Categorical(logits).sample gives an int index, with no unimix at HEAD; it is one-hot (bf16) for the RSSM.	p.18 (actor unimix, discrete only)	embodied/jax/heads.py:101-110, :146-155; embodied/jax/outs.py:161-186, :208-241; dreamerv3/configs.yaml:100-103; dreamerv3/main.py:249-258; embodied/core/wrappers.py:76-110; dreamerv3/rssm.py:137	conflict for discrete actions only (unimix). The continuous parameterisation is code_only.	2411f7d: discrete 'onehot' with unimix 0.01 (straight-through one-hot, argmax index sent to the env); continuous 'normal' with the same maths as bounded_normal (1a532a7 nets.py:486-493).
7.5 DMC proprio environment protocol	Paper protocol: 18 tasks (all 20 visual-control tasks minus quadruped walk/run), action repeat 2, 500K env steps (= 250K agent steps), 16 env instances, replay ratio 512 (0.5 updates per agent step, 0.25 per env step, about 125K updates in total), 12M model, about 0.3 A100 GPU-days. ActionRepeat sums rewards over the repeats and stops early on is_last or is_terminal; reset steps are not repeated. DMC episodes are 1000 env steps, i.e. 500 agent steps (+1 reset obs) at repeat 2. These tasks never terminate (is_terminal is always False). Observation: every key of the dm_control observation dict becomes its own vector input ('/' → '_'), float64 → float32 (UnifyDtypes), symlogged in the encoder. Action space [−1, 1]. Logger x-axis = agent steps × repeat. In proprio mode the env still renders a 64×64 'log/image' each step (logging only; skip it in Ajax). HEAD config dmc_proprio instead: <<size1m, env.dmc.image False, run {steps 1.1e6, train_ratio 1024}, env.dmc.repeat 1 (the default).	Table 2 p.19 (Proprio Control row: 18 tasks, 500K, repeat 2, 16 envs, ratio 512, 0.3 GPU days, 12M); p.19 "Benchmarks - Protocols" and "Computational choices" (control suites use 12M)	dreamerv3/configs.yaml:36 (dmc repeat 1), :178-182 (dmc_proprio); embodied/envs/dmc.py:21-76; embodied/core/wrappers.py:57-73; embodied/envs/from_dm.py:21-78; dreamerv3/main.py:155, :179	conflict: HEAD (1M params, RR 1024, repeat 1, 1.1M env steps) does not match the paper. 2411f7d does. To reproduce Table 11, use size12m, RR 512, repeat 2, 16 envs, 250K agent steps.	2411f7d dmc_proprio: <<size12m, run.train_ratio 512, run.steps 3e5 agent steps (= 6e5 env steps), env.dmc.repeat 2 (2411f7d configs.yaml:84, :233-237). This matches the paper; the same holds at f8817c4^. f8817c4 (2024-12-07) changed it to size1m / RR 1024 / steps 1.1e6 / repeat 1.
7.6 Per-benchmark run configs: paper Table 2 vs code	Format: paper env steps / repeat / envs / RR / model, then the HEAD config. Minecraft: 100M / 1 / 64 / 32 / 200M; HEAD sets only the task, so envs default to 16, RR to 32 and steps to 1e10. DMLab: 100M / 4 / 16 / 32 / 200M; HEAD steps 2.6e7 agent steps (= 104M frames), RR 32. ProcGen: 50M / 1 / 16 / 64 / 200M; HEAD steps 1.1e8 agent steps at repeat 1 (= 110M env steps), RR 64. Atari: 200M / 4 / 16 / 32 / 200M; HEAD steps 5.1e7 agent steps, RR 32. Atari100K: 400K / 4 / 1 / 128 / 200M; HEAD steps 1.1e5, envs 1, RR 256. BSuite: – / 1 / 1 / 1024 / 200M; HEAD envs 1, RR 1024. Proprio Control: 500K / 2 / 16 / 512 / 12M; HEAD size1m, steps 1.1e6, repeat 1, RR 1024. Visual Control: 1M / 2 / 16 / 512 / 12M; HEAD steps 1.1e6, repeat 1, RR 256, default 200M size. Crafter (Fig. 6 only): HEAD steps 1.1e6, envs 1, RR 512.	Table 2 p.19; p.19 "Environment instances" ("For Minecraft, we use 64 environments") and "Computational choices"	dreamerv3/configs.yaml:31, :36, :55 (run.envs 16), :155-196	conflict for Minecraft (env instances), ProcGen (budget), Atari100K (RR), Proprio and Visual control (size, RR, repeat, steps). Agree for DMLab, Atari and BSuite.	2411f7d: dmc_proprio size12m / RR 512 / repeat 2 / 3e5 agent steps; dmc_vision size12m / RR 512 / repeat 2 / 6e5 agent steps (both match the paper); atari100k RR 256 (already differed).
7.7 DMC proprio sanity targets	Dreamer at 500K env steps (5 seeds, paper Table 11): Acrobot Swingup 134, Ball In Cup Catch 962, Cartpole Balance 990, Cartpole Balance Sparse 990, Cartpole Swingup 852, Cartpole Swingup Sparse 491, Cheetah Run 614, Finger Spin 931, Finger Turn Easy 793, Finger Turn Hard 889, Hopper Hop 113, Hopper Stand 576, Pendulum Swingup 788, Reacher Easy 954, Reacher Hard 938, Walker Run 649, Walker Stand 964, Walker Walk 936. The printed "Task mean 871, Task median 754" rows are SWAPPED. Recomputing from the 18 rows gives mean 753.6 and median 870.5, and every column has the same swap. Corrected baseline mean/median: PPO 215/94, DDPG 689/771, DMPO 705/801, D4PG 733/792. Bundled 5-seed curves (scores/dmc_proprio-dreamerv3.json.gz, x = env steps 10K..490K): at 490K the task mean is 757.8 and the median 868.5; at 250K the mean is 675.4 and the median 790.4. High-variance tasks (population SD, ddof = 0; sample SD in parentheses): cartpole_swingup_sparse 393 (439), hopper_stand 225 (252), finger_turn_easy 203 (227). The scores are episode returns of the stochastic acting policy.	Table 11 p.34 (verified on the PDF); Figure 14 p.33	scores/dmc_proprio-dreamerv3.json.gz (added in 2411f7d)	agree (the bundled curves match the corrected summary)	2411f7d added this new file at a new path. The 2023 curves lived at scores/data/dmcproprio_dreamerv3.json.gz (423291a).

Implementation subtleties¶

Shared building blocks¶

BlockLinear init uses fan_in = the full input width I, not the per-block width I/g, because compute_fans multiplies the per-block fan by g. A per-block fan-in makes the GRU weights sqrt(8) ≈ 2.83× too large at init.
Every hidden Linear keeps its bias, even when RMSNorm follows. RMSNorm: eps 1e-4, statistics in f32, learnable scale only (no shift), output cast back to the activation dtype.
outscale applies only to the final head Linear; hidden layers always use 1.0. Policy 0.01 (mean and std Linears), reward 0.0, value 0.0, continue 1.0, RSSM logits 1.0, vector decoder 1.0 (HEAD) or 0.1 (2411f7d). Biases are always zero-initialised.
Two-hot: 255 bins built from a mirrored half for exact antisymmetry. Interpolation is linear in raw space between symexp-spaced bins, and targets are not symlogged. The expectation must use the symmetric pairwise sum, or the reward and value predictions are not exactly 0 at init. Targets must be f32.
Unimix (1%) is re-applied every time a categorical is built from logits, including inside both KL terms after the sg and when sampling. The straight-through estimator uses the mixed probabilities: sample = sg(onehot) + p_mix − sg(p_mix). Store raw logits.
DictConcat availability masking is all-or-nothing per key and per sample. Keys are sorted alphabetically. Discrete keys are one-hot encoded and never symlogged.
bf16 compute: Linear outputs, the RSSM carry and the GRU maths are bf16. Norm statistics, distributions, losses, returns and normalisers are f32. An all-f32 Ajax implementation changes numerics only. Note that if Ajax computes in f32, storing replay latents as bf16 is no longer lossless.

World model¶

Core: the RMSNorm after dynhid0 runs over the full D vector, not per block, and dyngru has no norm and no activation. GRU convention: deter' = u·cand + (1 − u)·deter, with u = sigmoid(u_pre − 1) and cand = tanh(sigmoid(r_pre)·cand_pre). The per-block gate layout is [r, c, u] within each block's 3D/g chunk; this matters only for official-weight compatibility.
is_first: zero deter, stoch and the previous action before the core, and zero the action again after one-hot encoding. h at an episode start is core(0, 0, 0). There is no learned initial state in the revised recipe. World-model losses at is_first steps are not masked.
KL: free bits clamp max(Σ_latents KL, 1.0) per (b, t), not per latent, and both dyn and rep are clamped. Both are KL(post ‖ prior); the sg goes on post for dyn and on prior for rep. The prior is computed from the same h_t that conditioned the posterior, in parallel after the scan, and h_t is not stop-gradiented, so dyn also trains the core and everything upstream (merge correction, 2.14).
Decoder: the input order is concat(stoch_flat, deter), while reward, continue, actor and critic use concat(deter, stoch_flat). The loss per key is the sum of squared errors over the feature dims against symlog(target), with no ½. Each obs key is its own loss term (scale 1), mean-reduced over (B, T).
Encoder: symlog only continuous keys. The tokens are the output of the 3rd Linear → RMSNorm → SiLU block, with no extra projection.
Action input to the dynamics is a / sg(max(1, abs(a))). The replay stores the unclipped policy sample; env clipping happens in a wrapper.
Stoch is sampled, not the mode, when acting, training and imagining.
Continue target = 0.997·(1 − is_terminal), a soft label. Truncations (is_last without is_terminal) are not zero targets. Never use is_last for this target.
Reward target = raw reward of the step that entered x_t (0 at is_first), not squashed.
Unused rssm.py class defaults (gelu, deter 4096, hidden 2048, classes 32) must not leak in. Use silu and the size-preset values.

Actor-critic¶

Targets and advantages bootstrap from the ONLINE critic (slowtar = False overrides the function default True). The slow critic is only a regulariser: two-hot cross-entropy toward its scalar mean, weight 1.0, in both the imagination and replay critic losses.
λ-return at step t uses r̂, ĉ and v of the NEXT state s_{t+1}, and r̂_0 is unused. Paper Eq. 5 literally indexes v_t; follow the code. lambda_return's val argument is never used.
The continuation weight w_t = Π_{i≤t} ĉ_i includes ĉ_0 of the start state. Slice weight[:, :−1] to align it with R_0..R_{H−1}. With contdisc the imagination recursion uses disc = 1. The replay recursion always uses γ = 1 − 1/333 with hard flags.
Return normaliser: update the EMAs of P5 and P95 first, then compute S = max(1, hi − lo) for the same batch. lo and hi start at 0 with no debias. Percentiles use linear interpolation over the flattened (B·K·H) = 15360 imagined returns, with no mask and no weights. Divide the advantage by S and do not subtract the offset.
Reinforce for both action types: log_prob of sg(sampled action). Continuous: Normal(tanh(μ), 0.9·sigmoid(s + 2) + 0.1), unsquashed samples and log-probs, no tanh-Jacobian term, entropy of the unsquashed Normal.
Imagined features are stop-gradiented (start feat via ac_grads = False, imagined feats always), so the actor and imagination critic never update the world model. The replay critic (repval) is not stop-gradiented and trains enc/dyn with scale 0.3. Its mask is ~is_last at the trained state, with no continuation product. Episode boundaries are cut by cont = λ(1 − is_last_{t+1}) and live = γ(1 − is_terminal_{t+1}).
Batch-major flattening of imagination starts (b major, t minor) is required for boot = R_imag[:, 0].reshape(B, K) to line up with the replay rows.
Reductions are means, not sums. Policy and value: mean over H, then over B·K. repval: mean over B·(T − 1). Weighted-out and masked entries still count in the denominator. The final states (imagination index H, replay index T − 1) receive no loss.
Normaliser EMA updates happen inside the loss forward pass and must be skipped in evaluation and report passes.
The policy samples stochastically in every mode; there is no deterministic eval action in the reference.

Optimizer and loss assembly¶

Optimizer order: AGC per tensor (0.3·max(1e-3, ‖p‖)) on the raw gradients, then RMS (β₂ 0.999, eps 1e-20 after the sqrt, bias-corrected), then momentum on the normalised update (β₁ 0.9, bias-corrected, its own step counter), then −lr. One optimizer covers all modules; the slow-critic EMA runs after the step.
The LR warms up linearly from 0 over 1000 updates, so the first update has lr = 0 (a no-op on the params, but the moments accumulate). This is code-only; the paper says "no annealing".
The AGC cap on zero-initialised tensors (gradient norm ≤ 3e-4 while ‖p‖ = 0) does not slow them: a constant clip factor cancels in the RMS normalisation.
Total loss = Σ scale·mean(term). Each vector obs key is its own rec term with weight 1. The loss must be an f32 scalar.

Replay, replay context and loop¶

Replay windows: an item at every start index (stride 1), length 65 = 64 trained steps + 1 context step. Windows cross episode and chunk boundaries. Do not restrict them to within-episode windows; the RSSM handles boundaries through is_first masking.
Batch annotation: is_first[:, 0] = True, and is_last = is_last OR (next step is_first), with the last column excluded. The repval mask relies on this.
Context alignment: the carry comes from the latent stored at index 0. The first trained obs is index 1. prevact for obs index t is the action stored at index t − 1 (data[act][:, 0:64]). Step 0 contributes no loss and is not overwritten by the write-back.
Stored per step: obs_t; a_t, the action chosen after seeing obs_t (zero if is_last_t); h_t and z_t, the posterior after filtering obs_t. After every update, write back the training-pass posteriors (pre-update params) for indices 1..64.
Online queue: per worker, push the item whose insertion happens at step n with n % 65 == 0 (starts 1, 66, 131, …). Train sampling pops it FIFO before uniform sampling; report sampling never pops.
The train-ratio clock counts agent transitions over all envs (not env frames). Updates per agent step = RR/(16·64). Evaluate after every transition. The first eligible call does exactly one update. Training starts once ≥ 1024 items exist (2048 transitions with 16 envs). There is no random prefill.
Capacity never binds for DMC proprio: 250K agent transitions ≪ 5e6 items, so FIFO eviction and the eviction retry never trigger there. A preallocated per-env ring of about 15.7K steps × 16 envs holds the entire run.
Memory for stored latents at 12M: D + 32·C = 2560 floats per step (10 KB in f32). Lossless compact storage is stoch as uint8 class indices [32] (C = 16 < 256) plus deter in bf16 (lossless only when compute is bf16), about 4.1 KB per step, or about 1 GB for 250K steps.
Functional, on-device replay design for Ajax (pure-JAX):
Per-env ring buffers [N_env, L_cap] with flags.
Valid starts are positions with 64 later steps in the same env stream.
Draw a uniform start index, gather 65 steps, and scatter the latent write-back at (env, start + 1 .. start + 64).
Emulate the online queue with a per-env pointer to the next unconsumed non-overlapping window (starts 1, 66, …), consumed before uniform draws.
Engineering staleness in the reference (acting params about two driver iterations stale, one-batch prefetch, outs and write-back delayed by one train call, replay RNG fixed at seed 0) is not algorithmic. A synchronous Ajax loop is a faithful simplification; document it as a deviation.
Scores: the reference reports training-episode returns of the stochastic policy. Paper Table 11's "Task mean" and "Task median" rows are swapped. Use mean ≈ 754 and median ≈ 871 at 500K env steps as the DMC-proprio sanity target, and mean ≈ 675 / median ≈ 790 at 250K env steps from the bundled curves.

Open questions¶

Each question gives the recommended default in the final table in bold. The user's call is needed where marked.

Resolved since: the user's answers are the fidelity rule of DESIGN.md §0 and the choices in deviations.md §1 and §3, which win over the recommendations below (notably 6 and 7, where Ajax follows 2411f7d).

DMC-proprio protocol (user call). The paper (Table 2) and 2411f7d use size12m, replay ratio 512, action repeat 2, 16 envs and 500K env steps. 2411f7d actually ran 3e5 agent steps = 600K env steps. HEAD (since f8817c4, Dec 2024) uses size1m, ratio 1024, repeat 1 and 1.1M env steps. This choice fixes every network width. Reproducing Table 11 needs the paper / 2411f7d setting. Recommended: paper (12M / 512 / repeat 2 / 500K).
LaProp β₂ (user call). The paper text says 0.99 ("its default parameters"); both code versions use 0.999, which produced the published curves. The LaProp reference's own defaults were not verified here (no literature fetch). Recommended: 0.999, with 0.99 as a config option.
Symlog MSE ½ factor. Paper Eq. 8 has ½; both code versions have none, so the code's reconstruction gradient is 2× the paper formula's. Recommended: code (no ½).
Continue head target. The paper uses a binary c_t with γ applied in the return; both code versions use the soft label 0.997·(1 − terminal) with γ folded in (contdisc = True). The returns are equivalent in expectation. Recommended: code (contdisc = True).
Discrete actor unimix. The paper (p.18, Table 4) and 2411f7d use 1%; HEAD's 'categorical' head drops it. This does not affect DMC, only discrete-action tasks. Recommended: 1% (paper), as a head option.
Start-state continuation and reward in imagination. HEAD uses the predicted ĉ_0 (≈ 0.997) in the weight. 2411f7d used the replay flag 1 − is_terminal (hard, without the 0.997 factor) and the replayed reward. The effect is a constant factor of about 0.997 on all actor-critic weights, plus different handling of terminal start states. The paper is silent. Recommended: HEAD.
Vector decoder output outscale. HEAD uses 1.0; 2411f7d (paper era) effectively used 0.1, a 10× init difference. The paper is silent. Recommended: HEAD 1.0, unless strict paper-era fidelity is preferred.
Table 3 12M recurrent units. The paper prints 1024, which contradicts its own 8d rule; both code versions use 2048. Recommended: 2048 (treat as a typo).
Replay context in Ajax (user call). The replay context needs a replay that stores per-step latents (2560 floats per step at 12M) and supports write-back by step id. If Ajax's on-device replay cannot provide this, the deviation must be chosen. The clean fallback is the official replay_context = 0 behaviour, which is a zero-state start (is_first forced at index 0; see 5.11), not a reuse of stale state. Burn-in is another option. Recommended: implement it (it is part of the paper).
Compute dtype (user call). The reference uses bf16 compute with f32 params, norms, distributions and losses. Should Ajax replicate bf16 (GPU-oriented) or default to f32 for CPU benchmarks and reproducibility? Recommended: a dtype option; f32 default for CPU tests, bf16 on GPU.
Asynchrony. Should Ajax reproduce the reference's acting-param lag (about two driver iterations), one-batch prefetch, one-call-delayed write-back and early first batch? The paper does not specify these. Recommended: synchronous, documented as a deviation.
Replay sampler seed. The reference fixes it to 0, independent of config.seed. Recommended: derive it from config.seed (documented deviation).
What Table 11 scores measure. The paper does not say whether they are training-episode returns of the stochastic policy (what the default train script logs) or separate evaluation episodes, and the code has no deterministic eval mode. Recommended: log training-episode returns as the primary metric, plus optional sampled-policy eval episodes.
Replay windows crossing episode boundaries. The code does this; the paper is silent. Recommended: follow the code.
elements.when.Ratio at HEAD comes from the external elements package, which is not available locally. Its semantics are taken from the 2411f7d code it was moved from, so a later change in that package cannot be ruled out. Low risk.
Online queue in a functional replay. The paper lists it as a component. A per-env "next non-overlapping window" pointer reproduces it. Its effect is small (about 0.2% of rows at RR 512 after the start-up batches). Recommended: implement it, behind a flag.

Hyperparameter defaults (revised / Nature recipe, vector observations, DMC-proprio paper protocol)¶

"Value" is the default recommended before implementation; deviations.md §1 wins where it chose otherwise (OQ 6, 7). "Source" gives where it comes from. "Alternative / note" gives the HEAD value when it differs, and other conflicts. Open-question numbers refer to the list above.

Name	Value	Source	Alternative / note
seed	0	configs.yaml:8	—
batch_size B	16	Table 4; configs.yaml:10	—
batch_length T (trained steps)	64	Table 4; configs.yaml:11	—
replay_context	1	p.18; configs.yaml:15	0 = zero-state start (5.11; OQ 9)
replay sequence length	65	main.py:187 (consec·T + context)	—
consec_train	1	configs.yaml:13	—
consec_report / report_length / report_batches	1 / 32 / 1	configs.yaml:12, :14, :58	diagnostics only
replay capacity	5e6 items	Table 4; configs.yaml:40	never reached for DMC proprio
replay chunksize	1024 steps	configs.yaml:46	storage detail
replay online queue	True (push every 65th insert per worker)	p.18; configs.yaml:41; replay.py:114	—
replay fracs uniform / priority / recency	1.0 / 0.0 / 0.0	configs.yaml:42; Table 4 caption	—
prio exponent / maxfrac / initial (disabled)	0.8 / 0.5 / inf	configs.yaml:43	unused
recexp (disabled)	1.0	configs.yaml:45	unused
train_ratio (replay ratio)	512	Table 2 (Proprio Control); 2411f7d configs.yaml:233-237	HEAD dmc_proprio 1024; HEAD default 32 (OQ 1)
batch_steps for the ratio	1024 (= B·T)	train.py:24-25	—
updates per agent step	0.5	derived (512/1024)	—
env instances	16	Table 2; configs.yaml:55	—
action repeat	2	Table 2; 2411f7d configs.yaml:84	HEAD 1 (OQ 1)
env steps (budget)	500K (= 250K agent steps, about 125K updates)	Table 2	2411f7d run.steps 3e5 agent steps (600K env steps); HEAD 1.1e6
DMC episode length	1000 env steps (500 agent steps + 1 reset obs)	dm_control; embodied/envs/dmc.py	—
training start gate	len(replay) ≥ 1024 items (2048 transitions with 16 envs)	train.py:70-71	—
random prefill / exploration steps	0	train.py	—
Prefetch amount	1	streams.py:34	reference only (6.3)
eval_envs / eval_eps	4 / 1	configs.yaml:56-57	train_eval script only
log_every / report_every / save_every	120 s / 300 s / 900 s	configs.yaml:52-54	infrastructure
env action range	[−1, 1] (ClipAction after NormalizeAction)	main.py:249-258	replay stores unclipped
model dimension d (MLP units, RSSM hidden)	256 (size12m)	Table 2/3; configs.yaml:125-128	HEAD dmc_proprio 64 (size1m) (OQ 1)
deter D (recurrent units)	2048 (= 8d)	Table 3 rule; configs.yaml:126	Table 3 prints 1024 (typo, OQ 8); HEAD dmc_proprio 512
stoch latents S	32	Table 3; configs.yaml:91	fixed at every size
classes per latent C	16 (= d/16)	Table 3; configs.yaml:126	HEAD dmc_proprio 4
GRU blocks g	8	p.18; configs.yaml:91	—
feature width (deter + S·C)	2560	derived	—
RSSM dynlayers / imglayers / obslayers	1 / 2 / 1	configs.yaml:91	code_only
RSSM absolute	False (posterior sees h_t)	configs.yaml:91	—
GRU update-gate bias offset	−1 (u = sigmoid(x − 1))	rssm.py:157	—
action input normaliser threshold	1 (a / max(1, abs(a)))	rssm.py:137	—
encoder MLP layers / units	3 / 256	p.18; configs.yaml:94	—
decoder MLP layers / units	3 / 256	p.18; configs.yaml:97	—
reward head layers / units	1 / 256	p.18; configs.yaml:98	—
continue head layers / units	1 / 256	p.18; configs.yaml:99	—
policy layers / units	3 / 256	p.18; configs.yaml:100	—
value layers / units	3 / 256	p.18; configs.yaml:101	—
CNN depth (unused for vector obs)	16 (= d/16)	configs.yaml:127	out of scope
latent unimix	0.01	Table 4; configs.yaml:91	—
actor unimix (discrete actions only)	0.01	Table 4; p.18; 2411f7d	HEAD effectively 0 (OQ 5)
free nats	1.0 (after summing over latents)	Table 4; configs.yaml:91	—
RMSNorm eps	1e-4	nets.py:364	—
init truncation bounds / std correction	[−2, 2] / 1.1368 × sqrt(1/fan_in)	nets.py:168-170	—
outscale RSSM obslogit / priorlogit	1.0	configs.yaml:91	—
outscale vector decoder heads	1.0	configs.yaml:97; rssm.py:304	2411f7d 0.1 (OQ 7)
outscale reward head	0.0	p.6; configs.yaml:98	—
outscale continue head	1.0	configs.yaml:99	—
outscale policy (mean, std or logits)	0.01	configs.yaml:100	—
outscale value head	0.0	p.6; configs.yaml:101	—
two-hot bins	255	configs.yaml:98, :101	—
two-hot range (symlog space)	[−20, 20] (±symexp(20) ≈ ±4.85e8 raw)	Eq. 10; heads.py:132-144	—
loss scale rec (per obs key)	1.0	Table 4 (β_pred); configs.yaml:86	—
loss scale rew	1.0	Table 4; configs.yaml:86	—
loss scale con	1.0	Table 4; configs.yaml:86	—
loss scale dyn	1.0	Table 4; configs.yaml:86	v1 0.5
loss scale rep	0.1	Table 4; configs.yaml:86	—
loss scale policy	1.0	Table 4; configs.yaml:86	—
loss scale value	1.0	Table 4; configs.yaml:86	—
loss scale repval	0.3	Table 4; configs.yaml:86	—
symlog-MSE factor	1 (no ½)	outs.py:129-141 (both code versions)	paper Eq. 8 uses ½ (OQ 3)
reward_grad / repval_loss / repval_grad / ac_grads	True / True / True / False	configs.yaml:88, :114-116	—
imag_last (K = T)	0 → K = 64, 1024 starts	configs.yaml:104	—
imagination horizon H	15 (16 states incl. start)	Table 4; configs.yaml:105	—
discount horizon 1/(1 − γ)	333 (γ = 0.996997)	Table 4; configs.yaml:106	—
contdisc (continue target factor)	True (0.997·(1 − terminal))	configs.yaml:107	paper: binary target + explicit γ (OQ 4)
λ (imagination return)	0.95	Table 4; configs.yaml:108	—
λ (replay return)	0.95	configs.yaml:109	—
actor entropy scale η	3e-4	Table 4; configs.yaml:108	—
slowreg (imagination and replay critic)	1.0	Table 4; configs.yaml:108-109	—
slowtar (both losses)	False (online critic for targets)	configs.yaml:108-109	function default True is overridden
slow critic EMA rate / every	0.02 (decay 0.98) / 1	Table 4; configs.yaml:110	—
retnorm rate / limit / perclo / perchi / debias	0.01 (decay 0.99) / 1.0 / 5 / 95 / False; lo = hi = 0 at init	Table 4; configs.yaml:111	—
valnorm / advnorm	none (offset 0, scale 1)	configs.yaml:112-113	—
policy minstd / maxstd / std pre-sigmoid offset	0.1 / 1.0 / +2.0	configs.yaml:100; heads.py:151	init σ ≈ 0.8927
learning rate	4e-5	Table 4; configs.yaml:87	—
LR warmup	1000 updates, linear from 0	configs.yaml:87	code_only
LR schedule / anneal	const / 0	configs.yaml:87	—
AGC clip / pmin	0.3 / 1e-3	p.18; Table 4; opt.py:109	—
LaProp eps	1e-20	Table 4; configs.yaml:87	—
LaProp β₁	0.9	p.18; configs.yaml:87	—
LaProp β₂	0.999	configs.yaml:87 (HEAD and 2411f7d)	paper text 0.99 (OQ 2)
weight decay	0.0	Table 4 caption; configs.yaml:87	—
compute dtype	bfloat16 (params f32)	configs.yaml:74	f32 option (OQ 10)
fp16 loss scale init / bounds / growth interval (inactive)	1e4 / [1e-4, 1e5] / 1000	opt.py:25-29, :83-92	only for float16
DMC-proprio sanity target at 500K env steps	task mean ≈ 754, median ≈ 871	Table 11 (rows swapped; recomputed)	bundled curves at 490K: 757.8 / 868.5; at 250K: 675.4 / 790.4