TD-MPC2: implementation reference (merged and verified)

This document gives the Ajax (JAX) implementer one place to work from instead of the paper and the code. Each row merges the extractor's spec with the adversarial verifier's corrections and the items the verifier found missing. Re-checks made during the merge are logged in §0.5.

---

0. Sources, pinned versions, conventions

0.1 Sources

Source	Pin	Notes
Paper	Hansen, Su, Wang, TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024, arXiv:2310.16828v2	https://arxiv.org/abs/2310.16828v2. Page numbers are PDF pages.
Code, HEAD	nicklashansen/tdmpc2 @ e9f5932 (2026-07-13, "fix q ensemble weight init")	https://github.com/nicklashansen/tdmpc2/tree/e9f59321933cbc8e11a002b842adc7d4ffae8ff1. Unprefixed line numbers refer to HEAD.
Code, paper era ("PE")	5f6fade (2024-03-27, "update pinned torchrl version")	Cited as @5f6fade:file:line.
Code, first commit	b67b21c (2023-10-25)	results/tdmpc2/*.csv (the paper's plotted curves) were added in this commit; 88095e7 only renamed them. A merge-time git diff b67b21c 5f6fade shows the agent algorithm is unchanged. The diff adds the mpc=false option and pixel support, deletes a dropout loop in track_q_grad that never ran (§2.16), and makes model_size optional. The replay buffer differs: b67b21c stores whole episodes and samples with RandomSampler + RandomCropTensorDict; 5f6fade uses SliceSampler. For fixed-length episodes the two sample from the same distribution.
Unmerged official branches	origin/uncertainty-regularization (tip 4c03df6, 2024-03-27); origin/vectorized_env (tip 97c1447, 2025-05-21)	Read for design reference only. Neither is on main.
DreamerV3	danijar/dreamerv3 @ e3f0224	Used only to cross-check a shared two-hot block.

0.2 Version lineage: commits that change behaviour

Commit(s)	Date	Change	Consequence for a port
b67b21c	2023-10-25	First commit. Buffer capacity = min(buffer_size, steps)//episode_length episodes (2000 DMC episodes = 1M transitions), RandomSampler + RandomCrop(H+1). track_q_grad contains a "set Q dropout p=0" loop that never runs.	Produced the paper's curves. The algorithm is the same as PE.
f313929	2023-11-25	Deletes the dropout loop that never ran.	None.
26c7211 / cc62c4c	2024-01-07	episodic-rl branch: first termination head (_terminated). Not an ancestor of 5f6fade.	HEAD-only feature (see add30b5).
54145a4, 0f3bc77, aa9c6f3	2023-12-27 to 2024-01-10	Switch to SliceSampler, with capacity counted in rows. A vendored sampler existed briefly, then torchrl-nightly was used.	Sampling is the same for fixed-length episodes.
5f6fade	2024-03-27	Paper-era pin.	Reference "PE" behaviour.
8b73181 / 603b67c	2024-09-25	Sibling commits "merged commits" (same parent 88095e7; both in HEAD lineage). torch.compile/tensordict refactor: the Q ensemble moves to from_modules (start of the Q-init regression); a detach-Q replaces track_q_grad (start of the task-embedding gradient leak); zero_grad moves after step (end of the PE stale-pi-grad clip quirk); Gumbel-max replaces np.random.choice.	See §0.3.
c1dd0c0	2024-10-27	Offline trainer: assert num_eps == capacity becomes a warning.	The PE offline path would have failed this assert on a complete dataset (§4.18).
32fc2bd, 804f9b3, db18653, e452ca7	2024-12-03 to 2024-12-25	Policy refactor: the tanh-Jacobian term drops out of the optimised entropy bonus, and ½ln2π is subtracted per dimension.	Algorithmic change (§2.12).
6117bc4	2024-12-10	DMC wrappers simplified; gymnasium 0.29.1; mujoco 2.3.1 → 3.1.2.	Small numeric differences only.
add30b5, 81eb170, 62be41a, eece801, 7942e90, 38b31a5, 7ec6bc8	2025-04-08 to 2025-05-02	Episodic support: (1-terminated) in TD targets, termination head and BCE loss, planning masking, all gated by cfg.episodic (default false).	HEAD-only. Not in the paper.
75212c3	2026-07-13	torchrl pinned to 0.8.1.	Sampler version.
e9f5932	2026-07-13	Each Q member initialised separately before stacking.	Restores PE init: HEAD == PE.

0.3 HEAD vs paper-era: algorithmic differences at a glance

Area	PE (5f6fade, same as b67b21c agent)	HEAD (e9f5932)	Paper	Suggested default for Ajax (needs user sign-off, see Open questions)
Policy entropy bonus	size·(Σ_d(−½ε²−logσ) − ½ln2π) − Σ_d log(relu(1−tanh(u)²)+1e−6). Gaussian part scaled by the number of valid action dims; tanh-Jacobian term unscaled and with gradient.	≈ −size·Σ_d(−½ε²−logσ−½ln2π). The tanh-Jacobian term has ~0 gradient.	"H … computed in closed form". No size scaling, no tanh.	PE (closest to the published results).
Task-embedding grads from the policy loss (MT)	None. track_q_grad(False) freezes the Q params and e.	Leak: written to e.grad, then applied by the next world-model step.	Gradients w.r.t. p only (Eq. 4).	Stop-gradient e (paper and PE).
World-model grad-clip norm	Includes the previous update's policy grads: ‖(g_wm, g_pi,prev)‖.	World-model grads only.	Not specified.	World-model grads only (HEAD).
Termination / episodic	Absent.	Optional (episodic=false default).	Absent (infinite-horizon MDPs).	Off, or an optional Extension.
Q-ensemble init	Each member N(0,0.02) weights, bias 0, last weight 0.	Same (restored by e9f5932).	Not specified.	Same.
Matmul precision	fp32 "highest".	TF32 ("high").	Not specified.	Explicit jax_default_matmul_precision.
ManiSkill2 termination	Never terminates.	May terminate early.	Fixed length, no termination.	PE (not in scope).
Final elite draw	np.random.choice(p=score)	Gumbel-max	Not specified	Equivalent. Ajax (M3): inverse CDF from one uniform, numpy's own choice algorithm, so the PE draw is replayable (3.19).
_prev_mean at start	Created lazily (the step-0 eval creates it)	Zeros buffer	Not specified	Zeros (equivalent in practice).
Replay buffer	SliceSampler (rows). b67b21c: episode storage.	SliceSampler (torchrl 0.8.1)	Uniform, capacity 1M	Episode ring buffer (= b67b21c).

0.4 Conventions

• Symbols: H = horizon = 3; B = batch; L = latent_dim; A = action dim (padded A_max in multi-task); K = num_bins = 101; N = num_samples = 512; P = num_pi_trajs = 24; E = num_elites = 64; Nq = num_q; T = episode length in agent steps (after action repeat); ST = single-task; MT = multi-task.
• Tensors are time-major (T, B, ·). ‖·‖ is the L2 norm. "Agent step" means one action-repeated env.step.
• Paper-vs-code labels: agree, conflict, code_only, paper_only.

0.5 Merge log: corrections applied and re-checks made during the merge

1. All verifier corrections are applied. No correction was rejected.
2. JAX truncated_normal. The model spec said it "truncates at ±2σ and rescales the std". The update verifier recalled a 1/0.8796 rescale. I checked the installed source (Ajax/.tox/py/.../jax/_src/nn/initializers.py, jax 0.10.2): random.truncated_normal(key, -2, 2, …) * stddev, with no std correction. With stddev 0.02 it truncates at ±0.04 and the realised std is ≈0.0176. Both claims are therefore off for this version. The conclusion is unchanged: use jax.nn.initializers.normal(0.02) (§1.13).
3. When the task-embedding leak started. The specs cite 8b73181 and 603b67c. Both are dated 2024-09-25, share parent 88095e7 and are in HEAD lineage, so both are cited.
4. Origin of the termination head. 26c7211 and cc62c4c are both "init" commits dated 2024-01-07 and both in HEAD lineage. Both are cited.
5. b67b21c vs 5f6fade. Verified with git diff: the agent algorithm is identical (see §0.1). The b67b21c buffer was confirmed in git show b67b21c:tdmpc2/common/buffer.py (capacity in episodes, RandomSampler, RandomCropTensorDict(H+1)).
6. The b67b21c dropout no-op. Confirmed. Ensemble.modules() was overridden to return the stateless nn.Sequential. Iterating it yields only the top-level NormedLinear/Linear children, none of which is an nn.Dropout, so m.p = 0 never ran.
7. Step units of the results CSVs. Confirmed: myo-hand-key-turn.csv (action repeat 1) is spaced 50,000 apart and walker-run.csv (action repeat 2) 100,000 apart.
8. 5f6fade update_pi was read directly: pi_optim.zero_grad at the start, track_q_grad(False), online Q(..., 'avg') in train mode, pi_loss = ((β·log_pis − qs).mean((1,2))·ρ).mean().
9. Task-embedding max_norm in JAX. The update spec recommended "project after each optimizer step". The model verifier showed this is not exact, and the training verifier showed the "stop-gradient scale" formula is self-contradictory. The lazy look-up-time renorm with write-back is adopted (§1.22).
10. Q-ensemble dropout at PE (correction made in M2 by running 5f6fade). Rows 1.1, 1.10, 2.3, 2.15, 2.16, 3.3 and 3.14, §2.A, the §2.C pseudocode and the §3.C pseudocode describe the HEAD behaviour; each carries a [PE: …] marker. At 5f6fade (and b67b21c) Ensemble wraps the functional module of combine_state_for_ensemble in a torch.vmap closure; that module is not a registered child, so model.train() / model.eval() never reach its nn.Dropout and Q dropout is active in every pass: the TD target's target-Q pass, the value and policy losses, and planning. HEAD's tensordict Ensemble registers self.module, so there dropout follows train/eval. See deviations.md §2.
11. Planning re-checked by running 5f6fade (M3). docs/world_models/parity/tdmpc2_plan_fixtures.py runs the unmodified act()/plan() with every random draw recorded and the per-iteration statistics read from plan()'s locals; Ajax's planner reproduces five chained decisions: t0, two warm starts, an eval-mode warm start, and t0 at an episode boundary while the agent holds a non-zero _prev_mean; with random reward, Q and policy heads, π trajectories are elites in some iteration of every decision (rows 3.5-3.21 confirmed). In the fixture the drawn elite is never a π trajectory (once the Gaussian samples converge they outrank the π trajectories), so that case of 3.19 is covered by Ajax-side tests. Corrections: (a) 3.23 and 3.15: with every value tied, torch.topk (CPU, torch 2.2.2) does not pick "arbitrary" elites but the first 64 candidates in index order, i.e. the 24 π trajectories and the first 40 Gaussian samples, as lax.top_k does (recorded; the CUDA order of the paper's runs is not verified). When tied values straddle the top-k boundary, its CPU topk selects a different subset of them than lax.top_k (lower index first), not only a different order (verified on torch 2.2.2 CPU in the M3 review; not recorded in the fixture). (b) 3.10 and §3.A: planning runs in eval mode but, at PE, with Q dropout (item 10); the model is verified to be in eval mode during plan(). (c) 3.10: the discount powers are Python floats (float64 products), rounded to float32 where they scale a tensor. (d) 3.19: np.random.choice(p=score) is inverse-CDF sampling with one uniform (cumsum, normalised by its last entry, searchsorted(side='right')); Ajax uses the same algorithm, in float32.
12. Multi-task re-checked by running 5f6fade (M8). docs/world_models/parity/tdmpc2_multitask_fixtures.py runs the unmodified agent with cfg.multitask=True on three synthetic tasks of different obs/action dims (padded), recording the embedding table on entry to update(), at the _td_target call, on entry to update_pi and on exit, and five act(..., task) decisions (all three tasks); Ajax reproduces every loss, gradient norm, table and planning statistic. Confirmed: (a) 1.22: embedding_renorm_ rescales only looked-up rows of norm > 1, by 1/(norm + 1e-7); inside one update it binds at the first look-up (encode(obs[1:]), pre-step) and again at the policy loss' first look-up after optim.step() (post-step: an Adam step can push a renormed row back above 1, which the fixture records), so "renorm once per update" would miss the second write-back (row 1.22 corrected); rows absent from a batch are not renormed and still move by Adam momentum; act() persists its renorm (Ajax: T15). (b) 1.9/1.23: masked dims are exactly 0 in the policy's action, and n in the entropy is the valid-dim count per sample. (c) 3.9/3.17: candidates, means and stds are masked as stated; the executed action's invalid dims are exactly 0 in both modes. (d) 2.20/3.12: the per-task discounts are a float32 tensor; in planning discount *= self.discount[task] multiplies float32 tensors (the single-task planner's powers are float64 Python floats). (e) The world-model clip includes _task_emb's gradient, and with track_q_grad(False) the policy loss gives it none (2.18).

---

1. World model architecture and math utilities

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
1.1 NormedLinear	y = act(LayerNorm(Dropout_p(x Wᵀ + b))). Order: Linear(bias=True) → Dropout (only if p>0; inverted dropout that scales kept units by 1/(1−p); active only in train mode [PE: dropout active in every Q pass; §0.5 item 10]) → LayerNorm(out_features) → activation. LayerNorm: last axis, eps = 1e-5, biased variance, learnable γ (init 1) and β (init 0). Default activation: Mish, x·tanh(softplus(x)). The only other activation used is SimNorm. Flax: nn.LayerNorm(epsilon=1e-5, use_fast_variance=False) (Flax defaults to 1e-6).	Sec 3.1 p.4 (LayerNorm + Mish); App A p.16 (1% Dropout after the first linear layer of each Q); App H p.26 listing; Table 8 p.27 (Activation: LayerNorm + Mish)	common/layers.py:94-111	agree	PE: Mish and Dropout were inplace=True. Numerically identical.
1.2 mlp builder	mlp(in, hidden:list, out, act=None, dropout=0): dims=[in]+hidden+[out]; for i in range(len(dims)-2): NormedLinear(dims[i], dims[i+1], dropout=dropout if i==0 else 0, act=Mish). Last layer: NormedLinear(dims[-2], out, act=act) (Linear→LN→act, no dropout) if act is given, otherwise a plain Linear(dims[-2], out) (no LN, no activation). Dropout only ever appears in the first hidden layer. Dynamics, reward, termination, π and each Q use hidden=[mlp_dim, mlp_dim], i.e. 3-layer MLPs.	App H p.26 ("all other components are 3-layer MLPs")	common/layers.py:121-133	agree	—
1.3 SimNorm	Input (..., D) with D divisible by V = simnorm_dim = 8. Reshape to (..., D/V, V), softmax over the last axis, reshape back. No temperature (τ=1). Groups: latent 128 gives 16, 512 gives 64, 768 gives 96, 1376 gives 172. It is the last activation of the encoder and the dynamics, placed after LayerNorm: Linear → LN → SimNorm. Paper self-contradiction: Eq. 5 defines softmax(z/τ) (τ→0 gives one-hot codes), while the App H prose says the reverse ("τ→∞ would force all mass onto single categories"). This does not matter at τ=1. If Ajax exposes τ, use Eq. 5 semantics (divide by τ) and note that the App H prose is reversed.	Eq. 5 p.4; Table 8 p.27 (V=8, τ=1); App H p.28-29 (listing + prose); Fig. 9 p.8 (multi-task: No Norm 46.8, SimNorm 51.0, LN+SimNorm 54.2 = default)	common/layers.py:74-91; config.yaml:64	agree	—
1.4 State encoder h(s, e)	mlp(obs_dim + task_dim, max(num_enc_layers−1, 1)·[enc_dim], latent_dim, act=SimNorm). ST input: the raw state vector, with no observation normalisation or symlog. MT input: concat([obs_zero_padded, e], −1). No dropout. 5M: NL(S+T→256, Mish) → NL(256→512, LN→SimNorm). 19M: S+T→1024→1024→768. 48M: S+T→1792×3→768. 317M: S+T→4096×4→1376. Must accept arbitrary leading axes: (B,S) for obs[0] and (H,B,S) for obs[1:]. RGB encoder (out of scope): ShiftAug pad 3, x/255−0.5, 4 convs (k 7/5/3/3, s 2/2/2/1, 32 channels, 64×64 input), ReLU after the first three convs only, Flatten → SimNorm with no LayerNorm.	Eq. 2 p.3; App H p.26 (encoder has 2–5 layers)	common/layers.py:136-164; common/world_model.py:25, 103-112	agree	—
1.5 Latent dynamics d(z, a, e)	z' = mlp(L + task_dim + A, [mlp_dim, mlp_dim], L, act=SimNorm)(concat([z, e, a], −1)). Concat order: z, then e (appended by task_emb), then a. Deterministic, no residual. The output passes Linear→LN→SimNorm, so predicted latents lie on the same product of simplices as encoder outputs.	Eq. 2 p.3; App H p.26 (in = 512+T+A; last act SimNorm)	common/world_model.py:26, 114-121	agree	—
1.6 Reward head R(z, a, e)	logits = mlp(L + task_dim + A, [mlp_dim, mlp_dim], max(K,1)=101)(concat([z, e, a])), with a plain last Linear. Scalar reward = two_hot_inv(logits). The last Linear weight is zero-initialised and its bias is 0, so the initial logits are 0 and the initial prediction is ≈0 (§1.19). Trained with soft_ce against raw rewards.	Eq. 2 p.3; Eq. 3 p.4; App H p.26	common/world_model.py:27, 32, 123-130	agree (the zero-init is code_only)	—
1.7 Termination head (HEAD only, episodic=true)	Instantiated only if cfg.episodic (default false): logit = mlp(L + task_dim, [mlp_dim, mlp_dim], 1)(concat([z, e])), termination(z) = sigmoid(logit). assert task is None, so it cannot be used in MT (this also breaks _update for MT+episodic). Init: standard N(0,0.02), bias 0, last layer not zeroed. Training and planning use are in §2.7 and §3.11.	none (the paper treats all tasks as infinite-horizon, App H p.28)	common/world_model.py:28, 132-141; tdmpc2.py:126-136, 257, 282-283, 294-297; config.yaml:7, 19	code_only	History (corrected): first written as _terminated on the episodic-rl branch (26c7211/cc62c4c, 2024-01-07; not an ancestor of 5f6fade). Entered HEAD lineage in merge add30b5 (2025-04-08) and was renamed _termination in 81eb170 (2025-04-08). From add30b5 until eece801 the BCE loss was computed unconditionally. 62be41a (2025-04-10): BCE-with-logits on zs[1:], coefficient 20 → 1. eece801 (2025-04-15): unnormalized flag and if cfg.episodic gate on the loss; 7942e90: gate in planning. 38b31a5/7ec6bc8 (2025-05-02): head instantiated only if episodic. The README says it is off by default for reproducibility.
1.8 Policy prior network and log-std bound	out = mlp(L + task_dim, [mlp_dim, mlp_dim], 2A)(concat([z, e])). Last Linear: N(0,0.02) weights, bias 0, not zero-init. mean, raw = split(out, 2, −1) (first A values are the mean). log_std = log_std_min + 0.5·(log_std_max − log_std_min)·(tanh(raw)+1) = −10 + 6·(tanh(raw)+1), range (−10, 2). log_std_min and log_std_dif are non-trainable buffers (YAML ints promoted to f32). At init (corrected): raw has std ≈ 0.02·√(512·0.454) ≈ 0.30 (E[mish(N(0,1))²] ≈ 0.454), so log_std ≈ −4 ± 1.8 (1σ) and σ ranges roughly 0.003 to 0.11 across dimensions and states. −4 (σ ≈ 0.018) is only the centre.	Eq. 2 p.3; Table 8 p.27 (log std −10/2); App H p.26 (Linear→2A)	common/world_model.py:29, 34-35, 154-155; common/math.py:12-13; config.yaml:45-46	agree for the bounds. The tanh rescaling is code_only.	—
1.9 Policy sampling, masking, log-prob, tanh squash	HEAD (world_model.py:144-184), on every call: mean, log_std as in 1.8; ε ~ N(0,I) with mean's shape (fresh each call). MT: m = action_masks[task]; mean*=m; log_std*=m; ε*=m; n = m.sum(−1, keepdim). ST: n = A. lp_g = Σ_{all A dims}(−½ε_i² − log_std_i − 0.9189385175704956) (keepdim → (...,1); each masked dim adds the constant −0.919). scaled = n·lp_g. u = mean + ε·exp(log_std); action = tanh(u); mean_out = tanh(mean); lp = lp_g − Σ_i log(relu(1 − action_i²) + 1e−6); entropy_scale = scaled/(lp + 1e−8) (not detached). Returns action, info{mean: tanh(mean), log_std: (masked) log_std, action_prob: 1., entropy: −lp, scaled_entropy: −lp·entropy_scale}. PE (@5f6fade:world_model.py:122-148, math.py): log_pi = n·(Σ_d(−½ε²−log_std) − ½ln2π) (constant subtracted once, then ×n; masked dims add 0). Then log_pi −= Σ_d log(relu(1−tanh(u)²)+1e−6) (unscaled). Returns (tanh(mean), tanh(u), log_pi, log_std). Masked dims give action = tanh(0) = 0 in both versions. Consumers: the policy loss (§2.11/2.12); the TD target and planner use the sample action; act(mpc=false, eval_mode) uses tanh(mean).	Eq. 4 p.4 (stochastic max-entropy prior); App A p.16 ("only compute policy entropy over valid action dimensions"); Sec 3.3 p.5 (masking prevents inflated entropy)	common/world_model.py:144-184; common/math.py:16-29	conflict. The Gaussian + tanh squash + Jacobian correction is code_only (the paper never mentions tanh or squashing). The single-sample log-prob and ×n scaling are code_only. The paper says closed form.	Changed in 32fc2bd (2024-12-03). Corrected attribution: 32fc2bd already introduced entropy_scale = β·scaled/log_prob with pi_loss = −(entropy_scale·entropy + qs). 804f9b3 added +1e-8, renamed the term to scaled_entropy, and moved β into update_pi. db18653 wrote scaled_entropy = −scaled_log_prob; e452ca7 reverted to −lp·entropy_scale six minutes later (same value). The gradient is the same in every HEAD-era variant. The return API changed from a 4-tuple to (action, info).
1.10 Q ensemble (architecture and vectorisation)	Nq members (5 by default; 2 at 1M; 8 at 317M). Each is mlp(L + task_dim + A, [mlp_dim, mlp_dim], 101, dropout=0.01): Linear(in→mlp) → Dropout(0.01) → LN → Mish → Linear(mlp→mlp) → LN → Mish → Linear(mlp→101). Parameters are stacked on a leading Nq axis. Forward: vmap(member, in_axes=(0, None), randomness='different')(params, x), so the input is broadcast and each member gets its own dropout mask. Output: (Nq, *batch, 101) logits. Dropout is on in train mode: the value-loss Q forward and the policy-loss Q forward. It is off for TD targets (target net forced to eval, and computed before model.train()) and for planning [PE: dropout active in every Q pass; §0.5 item 10]. JAX: split one dropout key per member and vmap over the keys.	Sec 3.1 p.4 (5 Qs); App A p.16; App H p.26; Table 8 p.27; Table 9 p.28	common/layers.py:8-33; common/world_model.py:30; tdmpc2.py:261-266, 280, 319	agree	PE: functorch combine_state_for_ensemble + torch.vmap(randomness='different') on an nn.ParameterList. 8b73181/603b67c: tensordict from_modules. Same forward, except that the PE ensemble's dropout ignores train/eval (§0.5 item 10).
1.11 Q() API	Q(z, a, task, return_type ∈ {min, avg, all}, target=False, detach=False). x = concat([z, e, a]) (e only in MT). net = target_Q if target else (detached_Q if detach else online_Q). out = net(x) has shape (Nq, ..., 101). 'all' returns logits. Otherwise idx = randperm(Nq)[:2]: one random distinct pair per call, shared by the whole batch and all time steps. Qs = two_hot_inv(out[idx]) has shape (2, ..., 1) (decoded scalars). 'min' → Qs.min(0); 'avg' → Qs.sum(0)/2. Min and avg act on decoded scalars, not on distributions. Use sites: TD target = Q(z', π(z').sample, target=True, 'min'); value loss = 'all' online; policy loss = Q(zs, π-action, detach=True, 'avg'); planner terminal = online 'avg'. Each call draws an independent pair. With Nq=2 the "random pair" is always both heads.	Sec 3.1 p.4 ("min of two randomly sub-sampled Q̄", Chen et al. 2021); Table 9 caption p.28 (always 2, regardless of Nq)	common/world_model.py:186-216; tdmpc2.py:136, 220, 255-257, 280-290	agree. The 'avg' for the policy loss and the planner is code_only.	PE: np.random.choice(Nq, 2, replace=False) (numpy RNG) and no detach flag (track_q_grad(False) instead). Same distribution.
1.12 Target Q, detached Q, modes	Detached Q = the online Q params' .data (shared storage, no grad). In JAX this is stop_gradient(online_q_params): gradients still flow to the inputs (z, a), never to the Q weights. Target Q = a clone of the online Q params made at construction, after the last-layer zeroing. Only the Q ensemble has a target: the encoder, dynamics, reward and π do not. The EMA update is in §2.15. WorldModel.train(mode) sets every submodule (including the detached Q) to mode, then forces the target Q to eval.	Sec 3.1 p.4 (Q̄ = EMA of each Q)	common/world_model.py:38-53, 69-86	agree	PE: target = deepcopy(Qs).requires_grad_(False), with no detach copy (track_q_grad).
1.13 Weight initialisation	(1) Every nn.Linear, NormedLinear included: trunc_normal_(std=0.02, a=−2, b=2). The bounds are absolute (±100σ), so in practice this is an untruncated N(0, 0.02²). Bias = 0. (2) LayerNorm: γ=1, β=0. (3) Task embedding: U(−0.02, 0.02). (4) Zero-init the reward head's final Linear weight and the final Linear weight of every Q member (params['2','weight']); their biases are already 0. (5) Each Q member draws independently (one key per member, vmapped init, then stacked). (6) The final layers of π, encoder, dynamics and termination are not zeroed. (7) Target Q is an exact copy of the online Q after (4). JAX: use jax.nn.initializers.normal(0.02). Do not use truncated_normal(0.02): in jax 0.10.2 it truncates at ±0.04 with no std correction (realised std ≈ 0.0176). See §0.5.	none	common/init.py:4-17; common/world_model.py:30-32, 41	code_only	PE == HEAD. Regression window 8b73181/603b67c (2024-09-25) to e9f5932 (2026-07-13): model.apply(weight_init) stopped reaching the from_modules-stacked params. Q members then kept PyTorch's default init (kaiming_uniform(a=√5), i.e. U(±1/√fan_in) for weights and biases). Only the final weight was zeroed, so the final bias stayed nonzero at ≈U(±1/√512) = ±0.044 and the initial Q logits were not uniform. Runs from that period differ.
1.14 symlog / symexp	symlog(x) = sign(x)·log(1+abs(x)); symexp(x) = sign(x)·(exp(abs(x))−1). The code uses log/exp; log1p/expm1 is numerically better and equal up to rounding. Applied only inside two_hot/two_hot_inv (reward and value targets and predictions), never to observations or latents.	Sec 3.1 p.4 ("log-transformed space"); App A p.16. The paper does not write out the formula (the code cites DreamerV3).	common/math.py:42-55	agree	PE: @torch.jit.script, same formulas.
1.15 two_hot (soft target)	Input x of shape (N,1). y = clamp(symlog(x), vmin=−10, vmax=10).squeeze(1); δ = bin_size = (vmax−vmin)/(K−1) = 0.2 (bins are linear in symlog space: −10, −9.8, …, 10); k = floor((y−vmin)/δ); o = (y−vmin)/δ − k; t[k] = 1−o; t[(k+1) % K] = o, with t in x's dtype. The wrap-around only happens at y == vmax, where o = 0, so it is harmless. Raw abs(x) > e^10−1 ≈ 22025 saturates at the edge bins. JAX: k = clip(floor(...), 0, K−2), o = (y−vmin)/δ − k. This gives identical targets and avoids an f32 edge case where (vmax−vmin)/δ rounds above 100.	Sec 3.1 p.4; Table 8 p.27 (101 bins)	common/math.py:58-71; common/parser.py:59; config.yaml:50-52	agree. vmin/vmax = ±10 and symlog-space placement are code_only.	PE: in-place scatter_, zeros without a dtype argument. Same math.
1.16 two_hot_inv (decode)	p = softmax(logits, −1); b = linspace(vmin, vmax, K) (symlog space, x.dtype); y = Σ_k p_k b_k (keepdim → (...,1)); return symexp(y). The expectation is taken in symlog space, then symexp is applied: symexp(E_p[b]) ≠ E_p[symexp(b)]. Used for reward predictions (planner) and for Q in 'min'/'avg'. The code sums naively, so a uniform p over symmetric bins gives ≈0, not exactly 0, in f32. DreamerV3 uses a symmetric summation. A shared TwoHot block should fix one summation order deliberately; the symmetric sum is harmless for TD-MPC2.	Sec 3.1 p.4	common/math.py:74-83	agree	PE cached bins in a global DREG_BINS. Same math.
1.17 soft_ce	soft_ce(pred (N,101), target (N,1)) = −Σ_k two_hot(target)_k · log_softmax(pred)_k (keepdim → (N,1)). Callers take .mean() over the batch, then weight by ρ^t and divide by H (and by Nq for values). The target carries no gradient (TD targets come from no_grad; rewards are data).	Eq. 3 p.4	common/math.py:5-9; tdmpc2.py:287-298	agree	—
1.18 num_bins ∈ {0,1} (do not port)	These modes do not work. The heads output max(K,1) = 1 logit, so log_softmax over a size-1 axis is identically 0 and the reward and value losses are exactly 0. K = 1 also divides by zero in bin_size. Expose only K ≥ 2, or implement regression explicitly as a separate head type.	—	common/math.py:5-9, 58-63; common/parser.py:59; common/world_model.py:27, 30	code_only	—
1.19 Initial reward/Q predictions	Zero final weights and zero biases make every initial reward/Q logit vector 0. two_hot_inv then gives symexp(mean(linspace(−10,10,101))) ≈ 0, identical for every input. Consequence: at init the planner sees equal values for all candidates (§3.23).	—	common/world_model.py:30-32; common/init.py:7-9; common/math.py:80-83	code_only	Not true in the 2024-09 to 2026-07 regression window (nonzero final Q bias).
1.20 Model-size table	MODEL_SIZE (enc_dim, mlp_dim, latent_dim, num_enc_layers, num_q): 1M 256, 384, 128, 2, 2; 5M 256, 512, 512, 2, 5; 19M 1024, 1024, 768, 3, 5; 48M 1792, 1792, 768, 4, 5; 317M 4096, 4096, 1376, 5, 8. num_q falls back to the config default 5 where the table omits it. config.yaml defaults equal the 5M values. ST online uses 5M (all released ST checkpoints expect model_size=5). Code quirk: task == 'mt30' and model_size == 19 sets latent_dim = 512 ("this checkpoint is slightly smaller"). model_size is ??? in config (HEAD: cfg.get('model_size', None); architecture args are free when it is unset).	Table 9 p.28; App H p.28 (no other changes across sizes); Table 8 p.27	common/__init__.py:1-24; common/parser.py:61-68; config.yaml:55-64	conflict: mt30-19M latent 512 (code) vs 768 (Table 9)	Same at 5f6fade. At b67b21c model_size defaulted to 5 and was mandatory.
1.21 Derived config	multitask = task ∈ {mt30, mt80}; tasks = TASK_SET[task] or [task]. task_dim: ST 0 (no embedding). MT: 96 if task == 'mt80' or model_size ∈ {1, 317}, else 64 (mt30 at 5/19/48M). bin_size = (vmax−vmin)/(K−1). MT always routes to the OfflineTrainer: there is no online MT training. cfg.action_dim = max over tasks; cfg.obs_shape = max obs dim; cfg.action_dims, obs_shapes and episode_lengths are per task.	Table 9 p.28 (task embedding 96 for all sizes); App H p.26 (Embedding(T, 96))	common/parser.py:56-78; envs/__init__.py:37-82; train.py:52	conflict (64 vs 96). The 64 exists only to load the released mt30 checkpoints. Use 96 for paper fidelity.	PE read cfg.model_size directly. HEAD uses cfg.get('model_size', 5) and converts the config to a dataclass for compile.
1.22 Task embedding e	W ∈ R^{n_tasks × task_dim}, nn.Embedding(n_tasks, task_dim, max_norm=1), init U(−0.02, 0.02) (‖w‖ ≈ 0.11 at d=96, so renorm is initially inactive). max_norm semantics: at every lookup (training, TD-target, act and planning calls included), each looked-up row with ‖w‖ > 1 is rescaled in place on the stored parameter, outside autograd: w ← w/(‖w‖+1e−7). The forward then gathers the (renormed) row, and the gradient is identity w.r.t. that row. Adam updates the renormed table; Adam moments are not touched by the renorm. Exact JAX recipe: before each lookup, rows = W[idx]; n = norm(rows); W = W.at[idx].set(where(n > 1, rows/(n+1e−7), rows)) outside grad, persisted in params; then e = W[task] (plain gather). Within one update this binds at two look-ups (verified in M8 by running 5f6fade, §0.5 item 12): the update's first, encode(obs[1:]) before the TD target, and the first after the world-model Adam step, pi(zs) in the policy loss; renorm and write back at both, for the batch's task ids, and renorm on act/plan look-ups. "Project all rows after each step" is only approximate: rows absent from a batch still move through Adam momentum, and torch gives R(u+Δ1+Δ2) where projection gives P(P(u+Δ1)+Δ2). With B=256 and 80 tasks ≈ 3 tasks miss a typical batch; with B=1024 this is rare. Concatenation (always appended after the main input): encoder [obs, e]; dynamics/reward/Q [z, e, a]; π/termination [z, e]. Broadcasting: 3-D (T,B,·) inputs with task (B,) repeat e over T; a single task id is repeated over the batch. task is int64: (B,) in training (task[0] of each slice) and (1,) in act/plan. Online, detached and target Q all use the same online table (no EMA of e). Optimised only by the world-model Adam (lr 3e-4) and included in its global clip.	Sec 3.3 p.5 (learnable, fixed-dim, ‖e‖ ≤ 1, conditions all five components); Table 8 p.27 (dim 96, norm 1); App F Fig 18-19 p.25 (normalised 54.2 vs unnormalised 46.6, 19M mt80); App H p.26	common/world_model.py:20-21, 88-101; common/init.py:10-11; tdmpc2.py:28	agree	Same at PE, except gradient routing in the π update (§2.18).
1.23 MT zero-padding and action masks (model side)	action_masks M ∈ {0,1}^{n_tasks × A_max}, M[i, :action_dims[i]] = 1 (valid dims are the leading dims; prefix masks), stored as a non-trainable buffer. Observations are zero-padded at the end to the max obs dim (MultitaskWrapper._pad_obs). The env receives action[:action_dims[i]]. π multiplies mean, log_std and ε by M[task] before sampling and log-prob (§1.9). Planner masking is in §3.9. Dynamics, reward and Q get the zero-padded action and do no masking internally.	Sec 3.3 p.5 ("zero-pad all model inputs and outputs … mask out invalid action dimensions … during both training and inference … only sample actions along valid dimensions during planning"); App A p.16	common/world_model.py:22-24, 158-164; tdmpc2.py:180-181, 195-197; envs/wrappers/multitask.py:17-57	agree	PE: a plain attribute moved in .to(). HEAD: register_buffer (8bbc14e handles its absence in ST checkpoint conversion).
1.24 Parameter-count unit tests	NormedLinear(i,o) has i·o + o + 2o params; Linear(i,o) has i·o + o. MT reference (S=39, task_dim=96, 80 tasks, A=6, 5M dims): encoder 167,936; dynamics 843,264; reward 631,397; π 582,668; Q (5 members) 3,156,985; task embedding 7,680; total 5,389,930. This exactly matches the paper. The paper overloads T: in_features use T = task_dim = 96, while Embedding(T, 96) uses T = 80 tasks (with T=80 in in_features the encoder would be 163,840). ST walker (S=24, A=6, no task): enc 139,520; dyn 794,112; rew 582,245; π 533,516; Q 2,911,225; total 4,960,618. ST dog (S=223, A=38): total 5,159,082. Target and detached Q copies are excluded from all counts.	App H p.26	common/world_model.py:20-30; common/layers.py:94-164	agree	—
1.25 Misc math utilities	gumbel_softmax_sample(p, temperature=1, dim=0): logits = log p, g = −log(Exponential(1)), argmax(softmax((logits+g)/T)). This is Gumbel-max, an exact categorical sample from p (temperature has no effect on the argmax). Used for the executed elite (score.squeeze(1), shape (E,)). termination_statistics(pred, target, eps=1e−9): termination rate and F1 at 0.5, logging only. int_to_one_hot is unused.	Sec 3.2 p.4-5	common/math.py:32-39, 86-109; tdmpc2.py:200, 329	code_only	PE: np.random.choice(p=score) on the CPU (same distribution).
1.26 dtype / precision	float32 everywhere; no AMP or bf16. HEAD torch.set_float32_matmul_precision('high') allows TF32 on Ampere+; PE uses full fp32. torch.compile(mode='reduce-overhead') for _update/_plan and Adam(capturable=True) are HEAD-only speedups that do not change the math. JAX on GPU defaults to reduced-precision f32 dots, so set jax_default_matmul_precision explicitly (use 'highest' in parity tests).	none	train.py:22-23; tdmpc2.py:30-31, 41-55	code_only	TF32 and compile are HEAD-only.

1.A Implementation subtleties (world model)

• Use NormedLinear order Linear → Dropout → LN → act. Dropout appears only in each Q member's first layer.
• Use LayerNorm eps 1e-5. Flax defaults to 1e-6.
• Initialise with normal(0.02) (torch's trunc bounds are absolute ±2), not truncated_normal(0.02). Zero the final weight of the reward head and of every Q member. Do not zero π's final layer. Init each Q member with its own key, then stack. Clone the target after zeroing.
• Concat order: [obs, e] for the encoder, [z, e, a] for dynamics, reward and Q, [z, e] for π and termination. The order is needed for checkpoint compatibility and the paper's parameter counts.
• SimNorm comes after LayerNorm in the last layer of the encoder and of the dynamics. It has no temperature.
• Q('min'|'avg'): one randperm(Nq)[:2] per call, shared by the whole batch and horizon. Reduce decoded scalars.
• two_hot: clamp in symlog space to [−10, 10]; 101 bins linear in symlog space; linear interpolation in symlog space. Decode = symexp(E_p[b]).
• Shared TwoHot block with DreamerV3: DreamerV3 (revised code, embodied/jax/heads.py:132-144, outs.py:273-330) puts 255 bins at symexp(linspace(−20, 0)) mirrored (raw space), interpolates two-hot weights in raw space, decodes as E_p[bins] with a symmetric summation, and uses log1p/expm1. A shared block therefore needs parameters for bin locations, interpolation space and decode transform. TD-MPC2's configuration is bins = linspace(−10, 10, 101) in symlog space with decode symexp(E_p[b]).
• Task embedding: use the lazy look-up-time renorm with write-back (§1.22), not differentiable normalisation. Target and detached Q use the online table.
• There is no observation normalisation and no symlog on observations or latents.
• MT masking: multiply mean, log_std and ε by the mask before computing the log-prob. Masked actions are then exactly 0.
• Use the parameter counts in §1.24 as architecture unit tests.

1.B Open questions (world model)

• Policy entropy formulation (§2.12): HEAD, PE or the paper's closed form? This is the most consequential choice.
• mt30 checkpoint compatibility: Table 9 (task_dim 96, latent 768 at 19M), or the code's 64/512? This only matters if loading released mt30 checkpoints is a goal.
• Termination head: include it as an optional Extension (ST only), or leave it out?
• Offline datasets: are invalid action dims in stored actions exactly 0? This was not verifiable because the datasets are external. Recommendation: the Ajax data pipeline zero-pads explicitly.
• SimNorm temperature: should Ajax expose τ at all (the paper is self-contradictory; the code has none)?

---

2. Training update (losses, targets, optimisers, EMA, RunningScale)

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
2.1 Batch layout and alignment	buffer.sample() returns obs (H+1,B,S), action (H,B,A), reward (H,B,1), terminated (H,B,1), task (B,) int64 or None. H=3, B=256 (1024 for MT). Storage row i holds obs_i plus the action, reward and terminated flag that led into obs_i; row 0 holds NaN placeholders. Each (H+1)-row slice yields action = td.action[1:], reward = td.reward[1:], terminated = td.terminated[1:] (or zeros_like(reward) if the key is absent: PE data and offline datasets), task = td.task[0]. Time-major via view(−1, H+1).permute(1,0). Transition t is (obs[t], action[t], reward[t], obs[t+1], terminated[t]), t = 0..H−1. Slices never cross episode boundaries, so no boundary masking is needed. Time-limit truncation bootstraps naturally.	Eq. 3 p.4 ((s,a,r,s')_{0:H} ~ B); Table 8 p.27 (uniform, batch 256); Sec 4.1 p.7 (batch 1024 MT)	common/buffer.py:17-25, 93-115; trainer/online_trainer.py:54-72, 104, 112	agree	PE returns (obs, action, reward, task) with no terminated. HEAD adds terminated (add30b5) and the zeros fallback (7ec6bc8).
2.2 Where terminated can be 1 (episodic)	A terminal flag sits on the last row of its episode, and slices never cross episodes. So inside a slice, terminated[t] can be 1 only at t = H−1: it gates only y[H−1], and no post-terminal latent, reward or value term ever appears. No extra masking of the consistency or reward losses is needed.	—	common/buffer.py:17-24, 100-106; trainer/online_trainer.py:101, 112	code_only	HEAD only.
2.3 TD target	Under no_grad, in eval mode (model.eval() was set at the end of the previous update), computed before model.train() [PE: dropout active in every Q pass; §0.5 item 10]. (1) next_z = enc(obs[1:], task), shape (H,B,L), with the online encoder (no target encoder). (2) a' ~ π(next_z, task): a fresh reparameterised sample (not the mean), tanh-squashed, masked in MT, from the pre-update π. (3) i,j = randperm(Nq)[:2]; Q̄_k = symexp(Σ_b softmax(target_logits_k)_b · bin_b); Qmin = min(Q̄_i, Q̄_j), shape (H,B,1). (4) y = r + γ·(1 − terminated)·Qmin. γ is a Python float in ST, or discount[task].unsqueeze(−1) with shape (B,1) in MT (broadcast over H). A one-step target for each t (no n-step, no λ-return). y is a raw-space scalar; two-hot/symlog happens inside the loss.	Eq. 3 p.4 (q_t = r_t + γ Q̄(z'_t, p(z'_t)), Q̄ = EMA); Sec 3.1 p.4 (min of two random Q̄); App A p.16; Table 9 caption p.28	tdmpc2.py:241-257, 261-263; common/world_model.py:186-216	conflict. Eq. 3 bootstraps from z'_t, the dynamics-predicted latent of Eq. 2. The code bootstraps from the encoded true next observation h(s_{t+1}) (the same tensor as the consistency target). Implement the code.	PE: y = r + γ·Qmin (no termination factor), same subsampling distribution, target-Q dropout on (§0.5 item 10). HEAD multiplies by (1−terminated) (add30b5).
2.4 Latent rollout and consistency loss	z_0 = enc(obs[0], task) with gradient, in train mode. For t = 0..H−1: z_{t+1} = d(z_t, action[t], task) (buffer actions, not policy actions); L_cons += ρ^t · F.mse_loss(z_{t+1}, next_z[t]), where F.mse_loss is the mean over batch and latent dim (B·L elements). next_z is a no_grad (stop-grad) target from the online encoder. Then L_cons /= H. ρ = 0.5, giving weights 1, .5, .25. Keep zs = [z_0 … z_H] of shape (H+1,B,L) (z_0 encoded, z_1..z_H predicted).	Eq. 3 p.4 (Σ_{t=0}^{H} λ^t ‖z'_t − sg(h(s'_t))‖²₂); Table 8 p.27 (coef 20, λ 0.5)	tdmpc2.py:261-276, 292; config.yaml:20-21, 39	conflict (3 ways): (1) the paper uses squared L2 (sum over the latent), the code uses the mean over the latent, a factor L = 512 smaller; coefficient 20 is tuned for the mean. (2) The paper has no /H. (3) The paper sums t = 0..H (H+1 terms), the code has H terms. Implement the code.	Same at PE.
2.5 Reward loss	r_logits = R(zs[0:H], action, task) with shape (H,B,K). L_rew = (1/H)·Σ_{t=0}^{H−1} ρ^t · mean_B soft_ce(r_logits[t], reward[t]), i.e. soft-CE against a two-hot target of symlog(r). Predictions at t ≥ 1 use predicted latents; gradients reach the reward head, dynamics, encoder and e.	Eq. 3 p.4 (CE(r̂_t, r_t)); Table 8 p.27 (coef 0.1)	tdmpc2.py:281, 287-288, 293; common/world_model.py:27, 32, 123-130	conflict: /H (code only), and the index range (paper's Eq. 3 sums t = 0..H, H+1 terms; code uses H terms).	Same at PE.
2.6 Value loss	q_logits = Q(zs[0:H], action, task, 'all') with shape (Nq,H,B,K), from the online ensemble in train mode (dropout on, independent mask per member). L_val = (1/(H·Nq)) Σ_t Σ_k ρ^t · mean_B soft_ce(q_logits[k,t], y[t]). Every head regresses the same y[t]. Gradients reach the Q params, e, and (through zs) the dynamics and encoder.	Eq. 3 p.4 (CE(q̂_t, q_t)); Table 8 p.27 (coef 0.1)	tdmpc2.py:280, 286-290, 298	conflict: /(H·Nq) (code only), and the index range t = 0..H−1 vs the paper's t = 0..H.	Same at PE.
2.7 Termination loss (HEAD, episodic only)	L_term = BCEWithLogits(Term(zs[1:H+1], unnormalized=True), terminated), the mean over all H·B elements, on the predicted latents z_1..z_H, with gradient into the dynamics and encoder. No ρ weighting and no /H beyond the mean. Coefficient 1. 0 when episodic=false.	none	tdmpc2.py:282-283, 294-297; config.yaml:19	code_only	Absent at PE. 62be41a changed the coefficient from 20 to 1 (2025-04-10).
2.8 Total world-model loss	L = 20·L_cons + 0.1·L_rew + 1·L_term + 0.1·L_val. Effective weights: 20/H·Σρ^t·MSE_mean; 0.1/H·Σρ^t·CE_r; 0.1/(H·Nq)·Σ_tΣ_kρ^t·CE_q.	Eq. 3 p.4; Table 8 p.27	tdmpc2.py:292-304; config.yaml:17-21	Coefficients agree. The normalisations are code_only (see 2.4-2.6).	PE has no termination term.
2.9 World-model optimiser	Adam with groups: encoder lr = 3e-4 × enc_lr_scale 0.3 = 9e-5; dynamics, reward, termination (if episodic, else empty), Q (online only) and task_emb (MT only) at lr 3e-4. betas (0.9, 0.999), eps 1e-8, weight_decay 0, no amsgrad. The policy is not in this optimiser; target and detached Q are not optimised. torch Adam matches optax.adam(eps=1e-8, eps_root=0). Adam is linear in lr, so the 0.3 scaling is equivalent to a per-group lr. HEAD: zero_grad(set_to_none=True) right after step(). optax: chain(clip_by_global_norm(20), multi_transform({encoder: adam(9e-5), rest: adam(3e-4)})).	Table 8 p.27 (Adam, 3e-4, encoder 1e-4)	tdmpc2.py:22-30, 307-310; config.yaml:22-23	conflict: encoder lr 1e-4 (paper) vs 9e-5 (code). Betas and eps are not in the paper.	PE: same groups minus termination, no capturable, optim.zero_grad at the start of update (see 2.10).
2.10 Gradient clipping	World model: clip_grad_norm_(model.parameters(), 20), one global L2 norm over all world-model grads together (encoder, dynamics, reward, [termination], Q, [e]), applied before the per-group lrs. Policy: clip_grad_norm_(π.parameters(), 20) separately. torch: coef = min(1, max_norm/(total_norm + 1e−6)); returns the pre-clip norm. optax clip_by_global_norm has no +1e-6 (negligible). HEAD: π grads are None during the world-model clip. PE (corrected): π grads are not cleared after pi_optim.step() (zero_grad runs at the start of update_pi), and model.parameters() includes _pi. So the world-model norm is sqrt(‖g_wm‖² + ‖g_pi,prev‖²), where g_pi,prev is the previous update's post-clip policy gradient. PE world-model grads are clipped harder and its logged grad_norm is inflated.	Table 8 p.27 (clip 20)	tdmpc2.py:228, 308; @5f6fade:tdmpc2.py:184, 236, 271	agree for the clip value. The PE inclusion of stale π grads is code_only.	Changed in 8b73181/603b67c, which moved π zero_grad to after step ("cudagraph prefers to see the zero grad after step").
2.11 Policy loss	Inputs: zs.detach(), shape (H+1,B,L) (z_0 encoded, z_1..z_H predicted, pre-step world-model params). a, info = π(zs, task) (reparameterised). Qp = Q(zs, a, task, 'avg', detach=True): the mean of two randomly drawn online heads with post-step params, param-stop-grad, dropout on. A new pair is drawn, independent of the TD-target pair. Gradient flows only through a (and log_std) into π. Then S.update(Qp[0]) and Qp = Qp/S (all H+1 steps, see 2.13). HEAD: L_π = mean_{t=0..H}[ρ^t · mean_{B,1}(−(β·scaled_entropy_t + Qp_t))], β = 1e-4, ρ^t = [1, .5, .25, .125]. The outer reduction is a mean over H+1 steps. PE: L_π = mean_t[ρ^t · mean_{B,1}(β·log_pi_t − Qp_t)]. Sign: Eq. 4 is printed as "maximise α Q − β H", which literally minimises entropy and contradicts "maximum entropy". It is a paper typo. Both code versions maximise Q/S + β·entropy_estimate. Follow the code sign.	Eq. 4 p.4 (Σ_{t=0}^{H} λ^t[αQ(z_t,p(z_t)) − βH], gradients w.r.t. p only, α via moving statistics); Table 8 p.27	tdmpc2.py:208-230, 313; @5f6fade:tdmpc2.py:184-197	conflict: (a) mean over t vs the paper's sum; (b) entropy form (2.12); (c) sign typo in Eq. 4; (d) HEAD MT task_emb leak (2.18). The H+1 latents, ρ weighting, α = 1/S and stop-grad on all but p agree.	PE: Q grads (and e) disabled by track_q_grad(False); online Q used directly in train mode.
2.12 Policy entropy bonus: HEAD vs PE	HEAD: scaled_entropy = −lp·scaled/(lp+1e−8), not detached, ≈ −n·lp_g in value and gradient (the extra term is −scaled·lp'·1e−8/(lp+1e−8)², negligible unless lp ≈ 0). The optimised bonus is β·n·(−lp_g), where lp_g is the pre-tanh Gaussian log-density of u. Its only reparameterised dependence is −log_std, so ∂bonus/∂log_std_i = β·n for valid dims. No tanh-Jacobian contribution, and the mean gets no entropy gradient. Hazard: if lp ≈ −1e−8 the denominator gives inf/NaN. If following HEAD, compute −n·lp_g directly (as db18653 briefly did). The constant offsets (HEAD: −n·A·0.919; PE: −n·0.919) only change the logged loss. PE: the loss term is +β·log_pi, with log_pi = n·(Σ_d(−½ε²−log_std) − 0.919) − Σ_d log(relu(1−tanh(u)²)+1e−6). The Gaussian part is ×n; the tanh correction is unscaled and has gradients w.r.t. mean and log_std. In both versions the Gaussian entropy is multiplied by the number of valid action dims, so the effective β is 1e−4 · #A_valid. info['entropy'] = −lp is for logging only.	Eq. 4 p.4 ("H … computed in closed form"); App A p.16 ("only compute policy entropy over valid action dimensions"); Table 8 p.27 (β = 1e-4). No action-dim scaling in the paper.	common/world_model.py:166-182; tdmpc2.py:226; @5f6fade:common/math.py (gaussian_logprob, squash)	conflict: closed form (paper) vs single-sample estimate (code); ×n scaling is code_only; HEAD drops the tanh-Jacobian.	Algorithmic HEAD vs PE difference, introduced in 32fc2bd (2024-12-03). Sign wording (corrected): in 32fc2bd the bonus inside pi_loss = −(entropy_scale·entropy + qs) is −β·scaled_log_prob, so the term's contribution to pi_loss is +β·scaled_log_prob, and the optimiser minimises β·n·log p_pre(u).
2.13 RunningScale (Q normalisation in the policy loss)	State: scalar S, init 1.0 (f32), percentiles [5, 95]. update(x) with x = Qp[0], shape (B,1): the t=0 Q values, unscaled, detached. Sort over the batch; pos = [5,95]·(B−1)/100; lo = floor(pos); hi = min(lo+1, B−1); w = pos − lo; p = sorted[lo]·(1−w) + sorted[hi]·w (the same as jnp.percentile(method='linear')). For B = 256, pos = 12.75 and 242.25. value = clamp(p95 − p5, min=1); S ← lerp(S, value, τ) = (1−τ)S + τ·value with τ = cfg.tau = 0.01 (the same hyperparameter as the target EMA). forward(x) = x/S. The update happens before dividing, so the current batch already influences S. S ≥ 1 always. MT uses one global S across tasks. Not saved in checkpoints (agent.save stores only model.state_dict()), so S resets to 1 on reload.	Sec 3.1 p.4 ("tuning α via moving statistics", citing DreamerV3); Table 8 p.27 (Moving (5%, 95%) percentiles). The paper does not give the EMA rate.	common/scale.py:5-47; tdmpc2.py:33, 221-222, 72-79	agree (the rate 0.01 is code_only)	PE: same math (read via .cpu().item(), written as x·(1/value)).
2.14 Policy optimiser	A separate Adam(π.parameters(), lr=3e-4, eps=**1e-5**, betas=(0.9, 0.999)). HEAD sequence: backward; clip(π, 20); step; zero_grad(set_to_none=True). One step per update, after the world-model step.	Table 8 p.27 (only Adam and 3e-4 globally)	tdmpc2.py:31, 227-230	code_only (eps 1e-5 and the separate optimiser)	PE: zero_grad at the start of update_pi. Not equivalent (corrected): stale π grads then enter the next world-model clip (2.10).
2.15 Target-network EMA	Only the Q ensemble has a target. The encoder, dynamics, reward and π do not. Init: an exact clone of the online Q after the last-layer zeroing. Update: θ̄ ← lerp(θ̄, θ, τ) = (1−τ)θ̄ + τθ, τ = 0.01 (momentum 0.99), over all Q leaves including LayerNorm γ/β. Runs once per update, after both the world-model step and the policy step, using post-step online params. The target is always in eval mode (no dropout) [PE: dropout active in every Q pass; §0.5 item 10].	Sec 3.1 p.4 (Q̄ = EMA of each Q); Table 8 p.27 (Q-fn momentum 0.99)	common/world_model.py:38-53, 74-86; tdmpc2.py:316; config.yaml:25	agree	PE: per-parameter lerp_ under no_grad. Same math.
2.16 Train/eval modes and dropout	Only Q has dropout (p = 0.01, first hidden layer, Linear → Dropout → LN → Mish). Within one update: (a) the TD target is computed in eval mode (target Q is always eval), so no dropout [PE: dropout on]; (b) model.train() precedes the rollout, so the value-loss Q forward has dropout on; (c) the policy update is still in train mode, so the detached-Q (HEAD) or online-Q (PE) forward has dropout on; (d) model.eval() at the end, so planning has no dropout [PE: dropout on]. Masks are independent per ensemble member (randomness='different') and per forward call. JAX: separate dropout keys for the value-loss and policy-loss Q forwards, and deterministic=True for the TD target and planning [PE: a dropout key for the TD target's target-Q pass too (M2), and one per MPPI iteration for planning's terminal Q pass (M3)].	App A p.16; Table 8 p.27 (1%)	common/layers.py:30, 94-133; common/world_model.py:30, 44-46, 74-80, 202-207; tdmpc2.py:32, 261-266, 319	agree	[PE: dropout active in every Q pass; §0.5 item 10] At PE the target-Q pass of the TD target and planning also have dropout (deviations.md §2). b67b21c's track_q_grad docstring says Q dropout is set to 0 during the π step, but the loop never ran (§0.5 item 6). Deleted in f313929. Effective behaviour: dropout on in the π step at every version.
2.17 Gradient-flow (stop-gradient) map	(1) next_z (consistency target and TD bootstrap state): no_grad. (2) y: no_grad, so no gradient to π or target Q. (3) The world-model loss gives grads to the encoder, dynamics, reward, Q, [termination], [e], and not to π. (4) The policy loss gives grads to π only: zs are detached and Q params stop-grad, but the gradient passes through the Q network w.r.t. its action input. (5) The termination loss uses zs[1:] with gradient into the dynamics and encoder.	Eq. 3 p.4 (sg on h(s'_t)); Eq. 4 p.4 (gradients w.r.t. p only)	tdmpc2.py:220, 261-263, 283, 313; common/world_model.py:40, 44-53	agree (except 2.18)	—
2.18 HEAD MT bug: policy gradient leaks into e	Since 8b73181/603b67c (2024-09-25), update_pi no longer disables e's grads. π(zs, task) and Q(…, detach=True) both look up _task_emb with requires_grad=True, so pi_loss.backward() writes ∂L_π/∂E into _task_emb.grad. pi_optim.zero_grad clears only π, and optim.zero_grad already ran (after the previous optim.step). The gradient therefore survives and is added to the next update's world-model gradient: included in that clip norm and applied by the world-model Adam at 3e-4. In effect, e is trained by the policy loss with a one-step delay. ST is unaffected. Recommendation: do not reproduce it. Apply stop_gradient(e) inside the policy loss (paper and PE).	Sec 3.1 Eq. 4 p.4 ("wrt. p only")	tdmpc2.py:219-230, 307-310; common/world_model.py:88-101, 150-151, 197-198	conflict (label corrected from code_only)	PE: track_q_grad(False) sets requires_grad=False on Q and e, so there is no leak. Present at every commit from 8b73181/603b67c to HEAD, including the pre-32fc2bd policy form.
2.19 Parameter freshness within one update	(a) The TD target uses the pre-step encoder, pre-step π and pre-EMA target. (b) zs fed to π are computed with the pre-step encoder and dynamics, then detached. (c) The Q heads (and MT e) inside the policy loss use post-step params: the detach view shares storage with the online params that Adam updated in place (PE uses the online Q directly, also post-step). (d) The target EMA uses the post-step online Q. (e) The next update's TD target uses the post-step π. Functional JAX: compute zs with old params → apply the world-model update → policy loss with new Q/e params and stop_gradient(zs_old) → update the target from the new params.	not specified	tdmpc2.py:261-316; common/world_model.py:40, 86	code_only	Same at PE.
2.20 Discount heuristic	γ(T) = min(max((T/5 − 1)/(T/5), 0.95), 0.995), T = env max_episode_steps after action repeat (discount_denom=5, min 0.95, max 0.995). DMC T=500 → 0.99; Meta-World/ManiSkill2/MyoSuite T=100 → 0.95; T=1000 → 0.995. ST: a Python float. MT: tensor discount[n_tasks] from cfg.episode_lengths; the TD target uses discount[task].unsqueeze(−1) (B,1); planning uses discount[task]. The same γ is used in TD targets and MPPI returns. For variable or unknown episode lengths the paper suggests the empirical mean length or γ = 0.99. HEAD gymnasium mujoco-* overrides: discount_max = 0.99, rho = 0.7 (episodic use case, HEAD only).	App H Eq. 7 p.28 (per-task γ in MT); Table 8 p.27 (heuristic)	tdmpc2.py:35-37, 57-70, 131, 256; config.yaml:26-28; envs/__init__.py:52-54, 81; envs/mujoco.py:57-58	agree	Same at PE (mujoco overrides are HEAD-only).
2.21 Order of operations	See §2.C pseudocode. The paper does not specify the ordering; reproduce the code ordering, including freshness (2.19).	—	tdmpc2.py:208-349 (_td_target 241-257, _update 259-332, update_pi 208-239, update 334-349)	code_only (label corrected from agree)	PE deltas are marked in the pseudocode.
2.22 Numerics / epsilons	float32 throughout. Squash 1e-6; entropy_scale denominator 1e-8 (HEAD); clip 1e-6; Adam 1e-8 (world model) and 1e-5 (π); LayerNorm 1e-5; embedding renorm 1e-7; MPPI 1e-9; termination_statistics 1e-9. HEAD TF32 matmuls; PE full fp32.	not specified	train.py:22-23; common/math.py:27; common/world_model.py:176	code_only	TF32 is HEAD-only.

2.C Reference pseudocode: one update() (HEAD; PE deltas in brackets)

obs[0..H], a[0..H-1], r[0..H-1], term[0..H-1], task = buffer.sample()        # time-major, slice-aligned (2.1)
# ---- targets: eval mode, no_grad ----                     [PE: target-Q dropout still on]
zn  = enc(obs[1:H+1], task)                               # (H,B,L), online encoder
a2  = pi_sample(zn, task)                                  # stochastic, tanh, masked (MT)
i,j = randperm(Nq)[:2]
y   = r + gamma[task] * (1 - term) * min(Qtgt_i(zn,a2), Qtgt_j(zn,a2))       # decoded scalars  [PE: no (1-term)]
# ---- world model: train mode (Q dropout on) ----
# [PE: optim.zero_grad() here]
z = enc(obs[0], task); zs = [z]; Lc = 0
for t in 0..H-1:
    z = dyn(z, a[t], task); Lc += rho^t * mean_{B,L}((z - zn[t])^2); zs.append(z)
ql = Q_all(zs[0:H], a, task)                               # (Nq,H,B,K), per-member dropout masks
rl = R(zs[0:H], a, task)                                   # (H,B,K)
Lr = sum_t rho^t * mean_B CE(rl[t], twohot(r[t]))
Lv = sum_t sum_k rho^t * mean_B CE(ql[k,t], twohot(y[t]))
Lc /= H; Lr /= H; Lv /= (H*Nq)
Lt = episodic ? mean BCEwithlogits(Term(zs[1:H+1]), term) : 0
L  = 20*Lc + 0.1*Lr + 1*Lt + 0.1*Lv
L.backward(); clip_global(all WM grads, 20)                # [PE: norm also includes stale pi grads]
Adam(enc 9e-5 | rest 3e-4, eps 1e-8).step(); zero_grad()
# ---- policy: still train mode ----                        [PE: track_q_grad(False) on Q AND task_emb]
zd = stop_grad(stack(zs))                                  # (H+1,B,L), PRE-step latents
ap, info = pi(zd, task)
k,l = randperm(Nq)[:2]; Qp = 0.5*(Q_k + Q_l)(zd, ap)        # POST-step params, param stop-grad, dropout on
S = lerp(S, max(p95(Qp[0]) - p5(Qp[0]), 1), 0.01); Qp = Qp / S
Lpi = mean_{t=0..H}( rho^t * mean_B( -(1e-4*info.scaled_entropy + Qp) ) )   # [PE: (1e-4*log_pi - Qp)]
Lpi.backward(); clip(pi grads, 20); Adam_pi(3e-4, eps 1e-5).step(); zero_grad(pi)   # [PE: zero_grad at start]
# ---- target EMA ----
Qtgt <- 0.99*Qtgt + 0.01*Q        # post-step, all Q leaves
model.eval()

2.A Implementation subtleties (update)

• The TD target bootstraps from the encoded next observation h(s_{t+1}), not from the dynamics prediction that Eq. 3 literally writes. It is a one-step target per t, with no λ-returns.
• Every Q head regresses the same y. The value loss is divided by H·Nq; the reward and consistency losses by H. F.mse_loss is a mean over batch and latent dim.
• ρ weighting: the world-model losses use t = 0..H−1 (1, .5, .25). The policy loss uses t = 0..H (1, .5, .25, .125) and takes a mean over H+1.
• Freshness: the policy loss uses pre-step zs and post-step Q (and e). The EMA uses the post-step Q. Make this explicit in a functional implementation.
• Dropout is on for the value-loss and policy-loss Q forwards and off for the TD target and planning, with independent masks per member and per call [PE: dropout active in every Q pass; §0.5 item 10].
• RunningScale: t=0 Q only, linear-interpolated 5/95 percentiles over the batch, clamp the range at ≥ 1, lerp with τ = 0.01, update before dividing, init 1. It is not checkpointed.
• The world-model clip is one global norm over all world-model params, not per group. π is clipped separately.
• Entropy: implement whichever version is chosen explicitly. Never compute HEAD's −lp·scaled/(lp+1e−8) literally; use −n·lp_g.
• MT: per-sample γ via discount[task]; prefix masks; one RunningScale shared by all tasks; stop-gradient e in the policy loss (recommended).
• In JAX, thread RNG keys for: π sample in the TD target, TD Q-pair permutation, value-loss dropout (Nq keys), π sample in the policy loss, policy-loss Q-pair permutation, and policy-loss dropout (Nq keys) [PE: and the TD target's target-Q dropout].

2.B Open questions (update)

• Entropy term: PE (|A|-scaled Gaussian + unscaled tanh correction), HEAD (|A|-scaled pre-tanh Gaussian, no tanh gradient), or the paper's closed form? None of the three agree. The 2024-12 change has no explanation in its commit message.
• TD bootstrap state: confirm h(s_{t+1}) (code) over the literal Eq. 3 z'_t.
• Loss normalisations (/H, /(H·Nq), mean over latent, H rather than H+1 terms): adopt the code form, since the coefficients are tuned for it?
• HEAD task-embedding leak: treat as a bug (stop-grad; recommended) or reproduce?
• PE stale-π-grad clip quirk: exclude (recommended, = HEAD) or reproduce for exact PE parity?
• Encoder lr: 1e-4 (paper) or 9e-5 (code)?
• Episodic termination support (HEAD-only): add as an optional Extension? Many gymnax envs terminate.

---

3. MPPI planning and action selection

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
3.1 Planning hyperparameters	mpc = true; H = 3; iterations = 6, +2 if action_dim ≥ 20 (8 for dog A=38 and humanoid A=21; mutated once in TDMPC2.__init__; MT uses the padded A_max = 6, so 6); N = 512 total candidates, of which P = 24 are policy-prior trajectories and 488 are Gaussian; E = 64 elites; min_std 0.05; max_std 2; temperature 0.5 multiplies: exp(0.5·(V − V_max)), an inverse temperature; no momentum.	Table 8 p.27; App A p.16 (momentum removed); Fig 5 caption p.6 (humanoid A=21; Table 2 p.18 says 24, which is a paper inconsistency; both give +2)	config.yaml:33-42; tdmpc2.py:34	agree. The paper does not say temperature multiplies.	Same at PE.
3.2 Shared horizon	cfg.horizon sets both the MPPI horizon (_estimate_value loop, π-trajectory rollout, mean/std shapes) and the training rollout length (zs is (H+1,B,L); losses summed over H). A modular port should either share one H (paper and code) or expose two hyperparameters and document the departure.	Table 8 p.27 (H listed under Planning); Eq. 3, Eq. 4 p.4	tdmpc2.py:127, 155-168 (planning); 269, 273-276, 292-298 (training)	agree	—
3.3 act() dispatch	act(obs, t0=False, eval_mode=False, task=None) under no_grad: obs → [1, S]; task → tensor([task]) with shape [1]. If mpc: plan(obs, t0, eval_mode, task) → action [A]. Otherwise use the policy prior (3.4). The model is always in eval mode while acting (Q dropout off) [PE: dropout active in every Q pass; §0.5 item 10].	Sec 3.2 p.4-5	tdmpc2.py:97-120, 45-55, 138-206	agree	PE encoded inside act() and passed z to plan. HEAD encodes inside _plan (for compile). No other algorithmic difference than Q dropout (§0.5 item 10).
3.4 act without MPC (Fig 9 "Policy" actor)	z = enc(obs); action, info = π(z). Training: the squashed sample. eval_mode: info['mean'] = tanh(μ) (masked in MT). Return action[0]. No clamp (tanh bounds it). Warm-start state is untouched. The "Planning" (no policy) ablation is num_pi_trajs = 0, supported by the if num_pi_trajs > 0 guards.	Fig 9 p.8 (Actor: Policy 42.2 / Planning 53.7 / Planning+policy 54.2)	tdmpc2.py:116-120, 154, 169	agree	PE: pi(z)[int(not eval_mode)][0]. Same semantics.
3.5 Policy-prior sampler inside planning	The π of §1.9, with fresh ε every call. Planning always uses the stochastic sample tanh(μ + ε·exp(log_std)), never tanh(μ), for both the 24 prior trajectories and the terminal action, even in eval_mode.	Eq. 2 p.3; Eq. 4 p.4; Sec 3.3 p.5	common/world_model.py:144-184	agree	PE: index [1] = sample.
3.6 Policy-prior trajectories	Once per decision step, before the MPPI loop: z = enc(obs) [1,L]; _z = repeat(z, 24); for t in 0..H−2: pi_a[t] = π(_z).sample; _z = d(_z, pi_a[t]); then pi_a[H−1] = π(_z).sample. pi_a has shape [H, 24, A]. These are independent stochastic rollouts through the deterministic dynamics, copied into actions[:, :24]. These columns are not resampled across iterations (only columns 24: are), but they are re-scored every iteration (new terminal π sample and new random Q pair).	Sec 3.2 p.5 ("a fraction of action sequences originate from p"); Table 8 p.27	tdmpc2.py:153-160, 168-170, 179	agree. Sample-vs-mean, fixed placement and no resampling are code-defined.	Same at PE.
3.7 Init and warm start (t0)	mean = zeros[H,A]; std = full[H,A](max_std = 2.0). If not t0: mean[:-1] = prev_mean[1:] (shift by one; the last row stays 0 and is not copied). std is never warm-started: it is reset to max_std at every decision step. If t0, the mean is all zeros (prev_mean is ignored, not cleared).	Sec 3.2 p.5 ("warm-start … initializing (μ, σ) as the solution to the previous decision step shifted by 1")	tdmpc2.py:164-167, 40	conflict: the paper warm-starts (μ, σ); the code warm-starts only μ.	PE: _prev_mean was created lazily. The first plan call had to be t0, or come after the step-0 eval (which in practice created it). HEAD: zeros buffer [H,A].
3.8 Gaussian candidates per iteration	ε ~ N(0,1) with shape [H, N−24, A] (fresh each iteration); cand = clip(mean[:,None] + std[:,None]·ε, −1, 1). The clamp comes before evaluation, so the evaluated and averaged actions are the clamped ones. actions[:, 24:] = cand. MT: actions *= mask[task] over the whole tensor, π columns included (they are already 0 on invalid dims).	Sec 3.2 p.5	tdmpc2.py:176-181	agree (the clamp is code_only)	Same at PE.
3.9 MT masking in planning	Applied (1) inside π; (2) to all candidates after sampling, every iteration; (3) to mean and std after the std clamp every iteration, so masked std = 0 < min_std, deliberately. In iteration 1 the invalid dims are sampled with std = 2 (the unmasked init) and then zeroed. Final noise is std[0]·ε with masked std 0, so invalid dims stay exactly 0. The warm-start mean is 0 on invalid dims. The env wrapper slices a[:a_i]. e is concatenated at every network call; the per-task γ is discount[task].	Sec 3.3 p.5	tdmpc2.py:180-181, 195-197; envs/wrappers/multitask.py:44-57	agree	Same at PE.
3.10 Trajectory value _estimate_value	Inputs: z [N,L] (the encoded obs repeated N times) and actions [H,N,A]. G = 0, disc = 1. For t in 0..H−1: r_t = two_hot_inv(R(z_t, a_t)) (the reward uses the current latent and action); z_{t+1} = d(z_t, a_t); G += disc·r_t; disc *= γ. Then a_H = π(z_H).sample (stochastic, fresh per candidate); return G + γ^H·Q_avg(z_H, a_H), where Q_avg is the mean of the decoded outputs of 2 random online heads. Output [N,1]. In formula: V = Σ_{t=0}^{H−1} γ^t r̂_t + γ^H · ½(Q_i + Q_j)(z_H, π(z_H)). All networks are online (no target), under no_grad, in eval mode [PE: the Q pass still has dropout; §0.5 item 10]. No RunningScale in planning. The discount powers are Python floats (discount *= self.discount, float64), rounded to float32 where they scale a tensor.	Sec 3.2 Eq. 6 p.4 (γ^H Q(z_{t+H}, a_{t+H}) + Σ_{h=t}^{H−1} γ^h R(z_h, a_h), expectation over (a_t..a_{t+H}) ~ N(μ, σ²), μ,σ ∈ R^{H×m})	tdmpc2.py:122-136; common/world_model.py:114-130, 186-216	conflict: Eq. 6 draws the terminal action from the planner Gaussian and writes H+1 actions while μ has H rows (internally inconsistent). The code uses π(z_H). Eq. 6's γ^h index is loosely written. The terminal 'avg' reduction is code_only.	PE (@5f6fade:tdmpc2.py:94-103): same formula, without termination masking, Q dropout on (§0.5 item 10).
3.11 Termination masking (HEAD, episodic only)	term = zeros[N,1]. Per t: r_t = R(z_t, a_t); z_{t+1} = d(z_t, a_t); G += disc·(1−term)·r_t (using term before this step); disc *= γ; if episodic: term = clip(term + float(sigmoid(Term(z_{t+1})) > 0.5), max=1) (hard 0.5 threshold, absorbing). Final: G + disc·(1−term)·Q_avg(z_H, π(z_H)). The reward of the transition into a predicted terminal state is kept; later rewards and the bootstrap are dropped. Reduces exactly to 3.10 when episodic=false. MT + episodic crashes (assert task is None).	none	tdmpc2.py:126, 130, 133-136; common/world_model.py:132-141	code_only	Absent at PE. History in §1.7.
3.12 Discount in planning	§2.20. ST: float; MT: discount[task] with shape [1], broadcast over [N,1]. Same γ as the TD target.	App H Eq. 7 p.28	tdmpc2.py:131-132	agree	Same at PE.
3.13 Decoding inside planning	two_hot_inv (§1.16): the expectation over 101 symlog-space bins, then symexp. Each Q head is decoded separately, then the scalars are averaged.	Sec 3.1 p.4	common/math.py:74-83; tdmpc2.py:128	agree	—
3.14 Q 'avg' and head subsampling	out = Q-ensemble(z, a) with shape [Nq, N, K] (online). qidx = randperm(Nq)[:2], drawn once per Q() call: the same pair for all 512 candidates within one iteration, a new pair each iteration. Returns (Q_i + Q_j)/2 of decoded scalars. Dropout is off (eval) [PE: dropout active in every Q pass; §0.5 item 10].	Sec 3.1 p.4; Table 9 p.28; App A p.16	common/world_model.py:186-216; common/layers.py:29-30	code_only (the paper only specifies min-of-two for TD targets)	PE: np.random.choice(Nq, 2, replace=False). Same distribution; dropout on (§0.5 item 10).
3.15 Elite selection	V = estimate_value(z, actions).nan_to_num(0) [N,1]: NaN → 0, ±inf → the largest/smallest finite f32 (torch default; jnp.nan_to_num matches). idx = topk(V[:,0], 64).indices; eV = V[idx] [64,1]; eA = actions[:, idx] [H,64,A]. Elites come from the current iteration's 512 candidates only, with no carry-over. Ties: lax.top_k puts lower indices first. torch's CPU topk agrees for a fully tied vector (the first 64 indices in order, §0.5 item 11), but when tied values straddle the top-k boundary it selects a different subset of them, not only a different order. Such ties come only from degenerate values (NaN → 0, ±inf, saturated decodes), where either choice is arbitrary; Ajax reproduces torch only for the fully tied vector (3.23).	Table 8 p.27; Sec 3.2 p.5	tdmpc2.py:184-186	agree	PE: in-place nan_to_num_.
3.16 MPPI score	score = exp(0.5·(eV − max(eV))); score /= score.sum(0), giving [64,1] that sums to 1. Values are raw return estimates (unnormalised), so the effective sharpness depends on the task's return scale.	Table 8 p.27 (temperature 0.5)	tdmpc2.py:189-191	agree (the formula is code-defined)	Same at PE.
3.17 Mean / std update	mean = Σ_k score_k·eA[:,k] / (Σ score + 1e−9) [H,A]. std = sqrt(Σ_k score_k·(eA[:,k] − mean)² / (Σ score + 1e−9)): weighted, biased (no Bessel), around the new mean. std = clip(std, 0.05, 2). MT: mean *= mask; std *= mask after the clamp. No momentum: the new (mean, std) fully replace the old ones.	Sec 3.2 p.5; App A p.16; Table 8 p.27	tdmpc2.py:192-197	agree	Same at PE.
3.18 Loop structure	For each of the 6 (or 8) iterations: sample columns 24: and clamp → MT mask → value + nan_to_num → top-64 → score → mean → std (around the new mean) → clamp std → MT mask mean/std. After the loop, score, eA, mean and std are those of the last iteration.	Sec 3.2 p.5	tdmpc2.py:173-197	agree	Same at PE.
3.19 Final action selection	k ~ Categorical(score_last) over the 64 elites (Gumbel-max: argmax(log score + Gumbel); JAX: jax.random.categorical(key, log(score)), where log 0 = −inf is fine). a = eA[0, k]: the first action of the sampled elite, not the mean. The chosen elite may be one of the 24 π trajectories.	Sec 3.2 p.5 ("executes the first action a_t ~ N(μ_t, σ_t)")	tdmpc2.py:200-202; common/math.py:86-94	conflict: the paper samples from N(μ_0, σ_0); the code uses a score-sampled elite plus σ_0 noise (3.20). μ is never executed directly.	PE: np.random.choice(arange(64), p=score), i.e. inverse-CDF sampling with one uniform u: cdf = cumsum(score)/sum, k = searchsorted(cdf, u, side='right') (float64; numpy RandomState.choice). Gumbel-max since 8b73181/603b67c. Same distribution. Ajax uses the PE algorithm (float32), so tests replay the reference's u.
3.20 Exploration noise, eval_mode, output clamp	s = std[0]: the last iteration's post-clamp std at step 0 (≥ 0.05 on valid dims; 0 on masked dims). If not eval_mode: a += s·N(0, I_A). Return clip(a, −1, 1). eval_mode only removes this noise. Evaluation stays stochastic: elite resampling, stochastic prior trajectories, stochastic terminal action and random Q pair.	Sec 3.2 p.5 (no eval variant defined)	tdmpc2.py:202-206	code_only	Same at PE.
3.21 Warm-start storage	After the loop: prev_mean ← mean, i.e. the last-iteration weighted mean [H,A], MT-masked. Not the selected elite and not the noisy executed action. Written on every plan call regardless of eval_mode. One prev_mean is shared by train and eval (the online trainer evaluates only at episode boundaries, and the next train call is t0, so they do not contaminate each other). Not checkpointed.	Sec 3.2 p.5	tdmpc2.py:40, 167, 205, 72-79	agree (μ only; see 3.7)	PE assigned before the noise; HEAD copies after. No behavioural difference.
3.22 Call sites, t0, seed phase	Online training: uniform random actions for step ≤ seed_steps. Planning from step seed_steps + 1, with t0 = (len(tds) == 1). Online eval: act(obs, t0=(t==0), eval_mode=True). Offline MT eval: act(…, eval_mode=True, task=i). evaluate.py: act(obs, t0=(t==0), task=i) with eval_mode left False (noisy). Stale warm start: on DMC, S = 2500 is an episode start, so step 2501 is the second step of an episode and the call has t0 = False. It warm-starts from the last mean of the step-0 evaluation (untrained model), not from zeros. 100-step domains behave the same way (S = 1000).	App H Eq. 8 p.28; Table 8 p.27	trainer/online_trainer.py:37, 76-88, 103-110, 115-123; trainer/offline_trainer.py:31; evaluate.py:80	agree	PE: same call pattern. PE eval_next = True at start vs HEAD False; both evaluate at step 0.
3.23 Planning at initialisation	With zero-initialised reward and Q heads (§1.19) every candidate has the same value ≈ 0 (exactly tied). top-k then returns the first 64 candidates in order on CPU torch 2.2.2 (verified by running 5f6fade, §0.5 item 11; the CUDA order is unverified): the 24 π trajectories and the first 40 Gaussian samples. The score is uniform, and mean/std are the plain mean and std of those 64 candidates. The step-0 eval runs in this state, and its final mean becomes the prev_mean inherited by the first post-seed planning call (3.22). Reproducing that call exactly therefore depends on matching the init.	—	common/world_model.py:30-32; common/math.py:80-83	code_only	—
3.24 Shapes, dtypes, numerics	One env per call: z [1,L] → [512,L]; actions [3,512,A]; V [512,1]; elites [3,64,A]; score [64,1]; mean/std [3,A]; output [A]. float32; 1e-9 in the mean/std denominators; everything under no_grad. HEAD: _plan is torch.compiled, with t0/eval_mode as Python bools (separate specialisations).	Sec 3.2 p.4-5	tdmpc2.py:111, 155-156, 163-170, 176, 185-186	agree	PE: no compile, no TF32.
3.25 MT planning runs only at evaluation	train.py selects the OfflineTrainer whenever MT, so there is no MT online acting loop. MT MPPI (masking, per-task γ, e) runs only in offline_trainer.eval (eval_mode=True) and in evaluate.py (noisy).	Sec 4.1 p.7	train.py:52; trainer/offline_trainer.py:31; evaluate.py:80	agree	—
3.26 Test-time uncertainty regularisation (App J, optional)	Paper: u_t = c · avg(q̂_1..q̂_N) · std(q̂_1..q̂_N) per latent. Objective γ^H(Q(z_{t+H}, a_{t+H}) − u_{t+H}) + Σ_h γ^h (R(z_h, a_h) − u_h). Table 10 (19M mt80): none 56.54; c=0.001 58.14; c=0.01 62.01; c=0.1 44.13. Not used in any main experiment. Unmerged branch: u(z) = mean(qs over ALL heads AND ALL 512 candidates) (a batch-wide scalar) × the per-candidate unbiased std over heads × uncertainty_coef, with qs = two_hot_inv(Q_all(z, π(z).sample)). Applied in the loop as G += disc·(r_t − u(z_{t+1})) (evaluated at the next latent) and at the terminal as Q_avg(z_H) − u(z_H) (so z_H is penalised twice). uncertainty_coef default 0.	App J Eq. 9-12, Table 10 p.30	origin/uncertainty-regularization:tdmpc2/tdmpc2.py:95-112; its config.yaml:41. Not on main.	conflict (branch vs Eq. 9: batch-wide mean; penalty at z_{t+1})	Branch merge-base 0f3bc77 (2024-01-08), tip 4c03df6 (2024-03-27), never merged.
3.27 Batched (multi-env) planning reference (unmerged)	origin/vectorized_env vectorises _plan over E envs: prev_mean [E,H,A]; candidates [E,H,N,A]; per-env top-k/gather; score normalised with +1e-9; one shared t0 for all envs (the trainer asserts that all envs reset together). Do not copy its final-index step (corrected): it passes log(score+1e−9) to a branch gumbel_softmax_sample that takes log(p+1e−9) again. Any score below 1 then yields a NaN Gumbel key, and the argmax (presumably NaN-as-max, not executed) returns elite 0, the top-valued elite, instead of a categorical sample. Use main's sampler on the raw score. Termination buffer (corrected): [N,1] broadcasts correctly against [E,N,1], so single-task episodic works on that branch. MT + episodic still hits the assert.	not in paper	origin/vectorized_env:tdmpc2/tdmpc2.py:41, 127-135, 167, 207-208; its common/math.py:89; trainer/online_trainer.py:42, 84	code_only	Tip 97c1447 (2025-05-21). Not on main.

3.C Reference pseudocode: plan() (HEAD, single env)

def plan(obs, t0, eval_mode, task):              # no_grad, eval mode (no Q dropout)  [PE: Q dropout on]
    z = enc(obs[None], task)                     # [1,L]
    if P > 0:                                    # P = 24
        _z = repeat(z, P)
        for t in 0..H-2: pi_a[t] = pi(_z).sample; _z = d(_z, pi_a[t])
        pi_a[H-1] = pi(_z).sample                # [H,P,A]
    z = repeat(z, N)                             # [N,L], N = 512
    mean = zeros[H,A]; std = full[H,A](2.0)      # std NEVER warm-started
    if not t0: mean[:-1] = prev_mean[1:]         # mean[-1] = 0
    actions[:, :P] = pi_a                        # fixed for all iterations
    for it in range(iterations):                 # 6 (+2 if A >= 20)
        actions[:, P:] = clip(mean[:,None] + std[:,None]*randn[H,N-P,A], -1, 1)
        if MT: actions *= mask[task]
        V  = nan_to_num(estimate_value(z, actions, task), 0)    # [N,1]
        ix = top_k(V[:,0], 64)
        eV, eA = V[ix], actions[:, ix]                          # [64,1], [H,64,A]
        score = exp(0.5*(eV - max(eV))); score /= sum(score)
        mean = sum_k score_k*eA[:,k] / (sum(score) + 1e-9)
        std  = sqrt(sum_k score_k*(eA[:,k] - mean)^2 / (sum(score) + 1e-9))
        std  = clip(std, 0.05, 2.0)
        if MT: mean *= mask[task]; std *= mask[task]
    k = categorical(log(score))                  # last iteration
    a = eA[0, k]
    if not eval_mode: a += std[0] * randn[A]
    prev_mean = mean
    return clip(a, -1, 1)

def estimate_value(z, actions, task):
    G, disc, term = 0, 1, zeros[N,1]
    for t in 0..H-1:
        r = two_hot_inv(R(z, actions[t]))
        z = d(z, actions[t])
        G += disc * (1 - term) * r
        disc *= gamma[task]
        if episodic: term = min(term + (sigmoid(Term(z)) > 0.5), 1)
    aH = pi(z).sample
    return G + disc * (1 - term) * Q_avg(z, aH)  # online Q, one random distinct pair per call, decoded then averaged

3.A Implementation subtleties (planning)

• N = 512 includes the 24 π trajectories: there are 488 Gaussian samples. π trajectories are generated once per step, never resampled, and re-scored every iteration.
• Prior trajectories and the terminal action use the stochastic squashed sample, even in eval_mode.
• The warm start shifts only the mean, and the last row is 0. std resets to 2 every step. t0 zeroes the mean. Store the final weighted mean, not the executed action.
• Clamp candidates to [−1, 1] before evaluation.
• In the value estimate, the reward comes from (z_t, a_t) before the dynamics step. The terminal value is γ^H · avg of 2 decoded online heads at π(z_H).sample.
• Score is exp(0.5·(V − max)) over elites. The 0.5 multiplies.
• std is the biased weighted std around the new mean, clamped to [0.05, 2]. MT masking comes after the clamp.
• There is no momentum and no elite carry-over.
• eval_mode removes only the final additive noise.
• Planning uses only online networks, in eval mode [PE: Q dropout still on, §0.5 item 10]. Target and detached Q are not used, and RunningScale is not applied.
• Functional/vmapped Ajax port (a design choice the reference does not cover): carry prev_mean [E,H,A] in the actor state; compute t0 per env from the previous done/reset flag; mean = where(t0[:,None,None], 0, shift(prev_mean)). Split keys for: the H π samples of the prior trajectories; per-iteration Gaussian noise, terminal-π ε and Q-pair permutation; the final categorical draw; and the exploration noise. Use lax.top_k and lax.fori_loop/scan (all shapes static).
• Seed phase: random actions while step ≤ seed_steps. The first planned call is usually t0 = False (stale warm start, 3.22).

3.B Open questions (planning)

• σ warm start: the paper warm-starts (μ, σ); the code warm-starts only μ. Recommendation: the code.
• Executed action: the paper uses a_t ~ N(μ*, σ*); the code uses a score-sampled elite plus σ*_0 noise. Which takes precedence under "as close to the papers as possible"?
• Eq. 6 terminal action: planner Gaussian (paper, internally inconsistent) or π(z_H) (code)? Confirm the code.
• Eval protocol: the trainers use eval_mode=True (still stochastic); evaluate.py is noisy. Which one produced the paper's numbers is unknown. Which should Ajax use?
• Vectorised envs: per-env t0 and prev_mean generalise the code, but neither the paper nor the code specifies them. Confirm the design.
• Stale warm start after seeding: reproduce exactly, or treat "no previous plan" as t0? The effect is negligible.
• App J uncertainty regulariser: in scope (multi-task offline)? It is off by default everywhere, and the branch formula differs from Eq. 9.
• Episodic termination masking in planning: include it as an Extension (ST only)?
• Separate planning vs training horizon: keep one shared H (the reference) or expose two?

---

4. Training loop, replay buffer, environments, multi-task machinery

Item	Spec (exact)	Paper ref	Code ref	Paper vs code	Version note
4.1 Entry point and trainer selection	cfg = parse_cfg(cfg); set_seed(cfg.seed) seeds python random, numpy, torch and torch.cuda (all devices). trainer_cls = OfflineTrainer if cfg.multitask else OnlineTrainer. Keyword evaluation order: env = make_env(cfg) (writes obs_shape, action_dim, episode_length and seed_steps onto cfg), agent = TDMPC2(cfg), buffer = Buffer(cfg), logger. MT is offline only and ST is online only. There is no online-MT path and no online finetuning of an MT model (evaluate.py warns that ST eval of MT checkpoints is unsupported).	Sec 3.1 p.4 (online alternation); Sec 4.1 p.7 (MT on fixed datasets); App J p.30 (MT is offline RL); App E p.24 (few-shot online finetuning)	train.py:46-60 (:52); common/parser.py:71; common/seed.py:7-12; evaluate.py:49-51	agree (finetuning is paper_only, 4.25)	Same at PE.
4.2 Online loop (exact control flow)	See §4.C pseudocode. Consequences: (i) eval happens at an iteration with step % eval_freq == 0, but only once the current episode has ended. On DMC (500 divides 50,000) this is exactly steps 0, 50k, 100k, …, including step 0 with an untrained agent. (ii) The episode still in progress at the end of the run is never stored or logged. (iii) The log x-axis is the agent step.	Sec 3.1 p.4; Table 8 p.27	trainer/online_trainer.py:74-127	agree (eval timing and episode-end insertion are code-defined)	PE eval_next = True initially vs HEAD False: no behavioural difference. HEAD adds the terminated field, the ValueError guard and extra metrics. b67b21c filled the row-0 action with torch.empty_like (garbage, never read).
4.3 Seeding phase	seed_steps S = max(1000, 5·T). DMC T=500 → S=2500; T=100 → S=1000; HEAD mujoco- T=1000 → 5000. Random actions for every step ≤ S (planning requires step > S): S+1 random actions. rand_act = action_space.sample() (f32, uniform on [−1,1]^A after action_scale). When updates begin (step == S), the buffer holds exactly the finished episodes: DMC episodes ending at steps 500…2500, i.e. 5 episodes = 2500 transitions (2505 rows); 100-step domains 10 episodes. Seed-phase actions are not reproducible from cfg.seed:* gym/gymnasium spaces carry their own lazily-seeded np_random, which set_seed does not touch. Env physics is seeded with task_kwargs={'random': seed}.	App H Eq. 8 p.28 (S = max(5T, 1000); called "environment steps" but counts agent steps); Table 8 p.27	envs/__init__.py:81-82; trainer/online_trainer.py:107-110; envs/wrappers/tensor.py:16-17; envs/dmcontrol.py:102-105	agree	Same at PE.
4.4 One-off pretraining burst	At the iteration where step == S (after that step's random env.step), agent.update(buffer) runs S times in a row (2500 on DMC, 1000 on 100-step domains). Each call samples a fresh batch and does the full update (world-model step, policy step, EMA). After that, exactly 1 update per agent step. Total updates = steps. JAX: lax.fori_loop/scan over S updates with split keys.	not described (Eq. 8 only defines S)	trainer/online_trainer.py:115-123	code_only	Present since b67b21c.
4.5 Update-to-data ratio	UTD = 1: after seeding, one update per agent step. One update = sample B = 256 slices of length H+1 = 4, then world-model step, policy step on zs.detach(), target EMA. There is no planning inside updates; TD targets use π and target Q.	Sec 4 p.6 (TD-MPC2 UTD 1, DreamerV3 512); Table 8 p.27	trainer/online_trainer.py:115-123; tdmpc2.py:259-349	agree	—
4.6 Per-step storage format (online)	An episode of T agent steps is stored as T+1 rows (DMC 501). Row k: obs = o_k (row 0 = reset obs); action = a_{k−1} (row 0: NaN); reward = r_{k−1} (f32 scalar; row 0 NaN); terminated = f32 0/1 flag of the transition into o_k (row 0 NaN; HEAD only); episode = int64 id, added in buffer.add. obs is f32 (a flat DMC state vector). There is no task key online. JAX: store zeros plus an implicit "row 0" convention instead of NaN, because a mask-by-multiply would propagate NaN into gradients.	Sec 3.1 p.4; Eq. 3 p.4	trainer/online_trainer.py:54-72, 104, 112; common/buffer.py:84-91; envs/wrappers/tensor.py:37-42	agree	PE: no terminated.
4.7 Replay capacity, storage, insertion	HEAD/PE: capacity = min(buffer_size=1e6, steps) rows (not transitions; a DMC episode uses 501). torchrl LazyTensorStorage, circular FIFO. Device heuristic: estimate bytes/row from the first episode and use cuda:0 if 2.5·bytes·capacity < free GPU mem, else CPU. add(td) extends by a whole episode. load(td) (offline) takes [N_eps, L], assigns episode ids arange(num_eps, num_eps+N) expanded over L, flattens to N·L rows and extends. b67b21c (the code that produced the paper's curves; corrected): capacity = min(buffer_size, steps)//episode_length episodes (DMC: 2000 episodes = 1M transitions), whole-episode storage, RandomSampler over episodes + RandomCropTensorDict(H+1), so there are no partial episodes. A JAX episode ring buffer [N_eps, T+1, ·] with N_eps = 1e6//T (2000 for DMC) therefore matches the original design; it is not a deviation.	Table 8 p.27 (capacity 1,000,000; uniform); App A p.16 (uniform instead of PER)	common/buffer.py:13-26, 50-67, 69-91; @b67b21c:common/buffer.py:40, 60-70, 95-101; config.yaml:29	agree	Row capacity and SliceSampler arrived in 54145a4 (2023-12-27) and moved to torchrl-nightly in aa9c6f3 (2024-01-10). PE: pin_memory=True, prefetch=1, offline episodes added one at a time. HEAD: pin_memory=False, prefetch=0, cache_values=cfg.multitask, bulk load().
4.8 Subsequence sampling	HEAD/PE: SliceSampler(num_slices=B, traj_key='episode', end_key=None, truncated_key=None, strict_length=True, cache_values=cfg.multitask). ReplayBuffer batch = B·(H+1) = 1024 rows → B slices of H+1 = 4 rows. Semantics from torchrl: trajectory boundaries are where episode changes; for each slice, a trajectory is drawn uniformly with replacement, then a start offset uniformly in {0, …, len−(H+1)}. Slices never cross episodes. A vendored copy in history (0f3bc77) confirms traj_idx = randint(n_traj, (B,)) and start = floor(U·(len − (H+1) + 1)). Corrected: in that copy strict_length=True raises RuntimeError if any stored trajectory is shorter than H+1, rather than dropping it, and boundaries are computed in linear storage order with the last index forced as an end (no circular-wrap handling). Behaviour at torchrl-nightly 2024.3.26 / 0.8.1 is unverified. For fixed-length episodes all variants are uniform over within-episode windows (DMC: starts 0..497 within 501 rows), which matches the paper. JAX recipe: ep = randint(0, n_valid_eps, (B,)), s = randint(0, T+1−(H+1)+1, (B,)), gather rows [s, s+H] via vmap(dynamic_slice). For variable-length data, store per-episode lengths and draw s ~ U{0..len_e−(H+1)}. With per-step writes from vectorised envs, sample only fully written windows.	Table 8 p.27; App A p.16; Eq. 3 p.4	common/buffer.py:17-25, 112-115; git show 0f3bc77:tdmpc2/common/samplers.py	agree	b67b21c: episode RandomSampler + RandomCrop (same distribution for fixed length). PE pinned torchrl-nightly 2024.3.26; HEAD torchrl 0.8.1 (75212c3).
4.9 Replay contents	Online, an episode's T+1 rows are added only when it ends, so transitions from the running episode are never sampled. HEAD/PE FIFO overwrite at capacity: a partially overwritten oldest episode remains as a shorter trajectory under its old id. With steps < 1e6 the row capacity equals steps but steps·(T+1)/T rows are written, so a few rows are overwritten near the end (negligible).	—	trainer/online_trainer.py:83-104 (:101); common/buffer.py:16-24, 66, 84-91	code_only	b67b21c: whole-episode FIFO (no partial episodes).
4.10 Batch preparation	§2.1. task comes from row 0 only, which is valid because a slice lies inside one single-task episode.	Eq. 3 p.4	common/buffer.py:93-115	agree	—
4.11 DMControl wrapper (state obs)	make_env: domain, task = task.replace('-','_').split('_',1); aliases cup→ball_in_cup, pointmass→point_mass; custom tasks registered under tag 'custom'. suite.load(domain, task, task_kwargs={'random': seed}, visualize_reward=False) → action_scale.Wrapper(−1, 1) → DMControlWrapper. step(a): cast a to the spec dtype, repeat 2 sim steps with the same action, reward = r1 + r2 (no break on last()), return the obs after the 2nd step, done=False, info = defaultdict(float) (success 0, terminated 0). Observation = np.concatenate of the flattened values of the dm_control observation OrderedDict in key order, f32, no normalisation. → Timeout(500): done = t ≥ 500 agent steps (= 1000 sim steps). → TensorWrapper (torch f32). Per-agent-step reward ∈ [0, 2]; max return 1000. Eval reuses the training env instance, so the env RNG continues across train and eval. Key order matters only for loading released checkpoints.	App C Table 6 p.20 (DMC 1000 / AR 2 / 500); App C p.20 (fixed length, no termination)	envs/dmcontrol.py:47-60, 92-111; envs/wrappers/timeout.py:4-25; envs/wrappers/tensor.py:26-42	agree	PE: ActionDTypeWrapper → ActionRepeatWrapper(2) (reward += r·discount, break on last()) → action_scale → ExtendedTimeStep → TimeStepToGym. Equivalent on DMC. 6117bc4 (2024-12-10) simplified this and upgraded the physics (mujoco 2.3.1 → 3.1.2, dm-control 1.0.16): small numeric differences. RGB path (3 stacked 64×64 frames) out of scope.
4.12 Time limits vs terminations	All paper tasks are infinite-horizon with fixed length. The Timeout done is not stored as terminal; terminated comes only from env info (always 0 on DMC/MW/MyoSuite, and on ManiSkill2 at PE). The last transition bootstraps from the stored final obs. Planning returns also ignore episode ends. HEAD episodic mode is described in §1.7, §2.7 and §3.11. The online trainer raises ValueError if a termination occurs while episodic=false.	App C p.20; Sec 2 p.2-3; Eq. 3 p.4	envs/wrappers/timeout.py:21-25; tdmpc2.py:242-257, 123-136; trainer/online_trainer.py:91-93	agree (with episodic=false)	Episodic support is HEAD-only (2025-04/05).
4.13 Episode length and action repeat by domain	Agent-step lengths (after action repeat): DMControl 500 (AR 2). Meta-World 100 (AR 2, 200 sim steps; goal-observable v2, _freeze_rand_vec=False; reset() takes one zero-action step and returns the pre-step obs, a code quirk). ManiSkill2 100 (AR 2). MyoSuite 100 (AR 1, success = info['solved']). HEAD-only gymnasium mujoco- (Walker2d-v4, HalfCheetah-v4): 1000 steps, AR 1, overrides discount_max=0.99, rho=0.7. The success metric is the success flag of the final step. ManiSkill2 (corrected): breaking the repeat early on done and setting info['terminated'] = done are HEAD-only*. At PE the wrapper discarded the env done and always returned done=False, which matches App C. At HEAD, ManiSkill2 can end episodes early (conflicting with App C), and with episodic=false such a done hits the ValueError guard. Use PE semantics.	App C Table 6 p.20 (DMC 1000/2/500; MW 200/2/100; MS2 200/2/100; Myo 100/1/100; total env steps DMC 4M-14M, MW 2M, MS2 4M-14M, Myo 2M); App C p.20 (final-step success)	envs/dmcontrol.py:57-60, 110; envs/metaworld.py:17-28, 50; envs/maniskill.py:47-55, 80; envs/myosuite.py:30-33, 56; envs/mujoco.py:53-58; trainer/online_trainer.py:44	agree at PE; ManiSkill2 at HEAD is a conflict	HEAD Meta-World/ManiSkill2/MyoSuite make_env assign env.max_episode_steps on HEAD's Timeout, where it is a read-only property. This appears to raise AttributeError at HEAD. PE used gym TimeLimit. Only DMC is in scope for vector-obs ST; mt80 needs Meta-World (obs 39, act 4).
4.14 Discount and iterations heuristics	§2.20 and §3.1.	App H Eq. 7 p.28; Table 8 p.27	tdmpc2.py:34-37	agree	—
4.15 Task sets mt30 / mt80 (order = task id = embedding row)	mt30: walker-stand, walker-walk, walker-run, cheetah-run, reacher-easy, reacher-hard, acrobot-swingup, pendulum-swingup, cartpole-balance, cartpole-balance-sparse, cartpole-swingup, cartpole-swingup-sparse, cup-catch, finger-spin, finger-turn-easy, finger-turn-hard, fish-swim, hopper-stand, hopper-hop (19 original), then walker-walk-backwards, walker-run-backwards, cheetah-run-backwards, cheetah-run-front, cheetah-run-back, cheetah-jump, hopper-hop-backwards, reacher-three-easy, reacher-three-hard, cup-spin, pendulum-spin (11 custom). mt80 = mt30 (ids 0-29, same order) + 50 Meta-World tasks: mw-assembly, mw-basketball, mw-button-press-topdown, mw-button-press-topdown-wall, mw-button-press, mw-button-press-wall, mw-coffee-button, mw-coffee-pull, mw-coffee-push, mw-dial-turn, mw-disassemble, mw-door-open, mw-door-close, mw-drawer-close, mw-drawer-open, mw-faucet-open, mw-faucet-close, mw-hammer, mw-handle-press-side, mw-handle-press, mw-handle-pull-side, mw-handle-pull, mw-lever-pull, mw-peg-insert-side, mw-peg-unplug-side, mw-pick-out-of-hole, mw-pick-place, mw-pick-place-wall, mw-plate-slide, mw-plate-slide-side, mw-plate-slide-back, mw-plate-slide-back-side, mw-push-back, mw-push, mw-push-wall, mw-reach, mw-reach-wall, mw-shelf-place, mw-soccer, mw-stick-push, mw-stick-pull, mw-sweep-into, mw-sweep, mw-window-open, mw-window-close, mw-bin-picking, mw-box-close, mw-door-lock, mw-door-unlock, mw-hand-insert. Padded dims: mt30 obs 24, act 6; mt80 obs 39, act 6.	Sec 4.1 p.7 (80 = 50 MW + 30 DMC = 19 original + 11 new); App C Tables 2-3 p.18-19	common/__init__.py:26-60; common/parser.py:78	agree	Same at PE.
4.16 MultitaskWrapper (padding, action slicing)	obs_dims, action_dims and episode_lengths per sub-env. observation_space = Box(−inf, inf, (max obs,), f32); action_space = Box(−1, 1, (max act,), f32). reset(task_idx=−1) selects envs[task_idx] (the default −1 is the last task) and returns the obs zero-padded at the end. step(a) passes the prefix a[:a_i] to the sub-env and returns the padded obs. Masks are prefix masks (§1.23). Each sub-env is a full ST stack TensorWrapper(Timeout(domain wrapper)). cfg.episode_length and seed_steps come from envs[0] via attribute forwarding; the offline trainer does not use them.	Sec 3.3 p.5; App A p.16	envs/wrappers/multitask.py:11-27, 44-57; envs/__init__.py:37-55	agree	Same at PE.
4.17 Learnable task embedding	§1.22.	Sec 3.3 p.5	—	agree	—
4.18 Offline MT datasets	Files: data_dir/*.pt (sorted glob). mt30 = 4 chunks (chunk_0..3), mt80 = 20 chunks (chunk_0..19), from huggingface.co/datasets/nicklashansen/tdmpc2/{mt30,mt80}/chunk_i.pt. Each is a TensorDict [N_eps, L] with L asserted == 501 (mt30) or 101 (mt80). Keys inferred from consumers: obs (zero-padded, f32), action (zero-padded), reward, task (int id per row). terminated is absent, so it defaults to zeros; episode is assigned at load. Buffer capacity: mt30 345,690,000 rows (690,000 episodes × 501 = 345M transitions); mt80 550,450,000 rows (5,450,000 × 101 = 545M transitions). A warning is printed if the episode or file count differs. The README says 128 GB RAM for mt80. mt80 composition (inferred, not proven): 5,450,000 = 345M/100 (3,450,000 DMC 100-transition segments, i.e. mt30's DMC data cut into 101-row segments) + 200M/100 (2,000,000 MW episodes). The per-task γ for DMC tasks in mt80 still comes from the env episode length (500 → 0.99).	Sec 4.1 p.7 (545M transitions from 240 single-task agents' replay buffers, 80 tasks; 345M for the 30-task subset; random to expert)	trainer/offline_trainer.py:42-65; common/buffer.py:69-82, 102-109; datasets/download_mt30.sh; datasets/download_mt80.sh	agree	Corrected: the 5f6fade offline path cannot run on a complete dataset. It asserts num_eps == capacity, but at 5f6fade capacity counts rows (550,450,000) while num_eps counts episodes (5,450,000). The assert was valid only at b67b21c (capacity in episodes). c1dd0c0 (2024-10-27) turned it into a warning. So the paper's MT models came from b67b21c-era code with the episode buffer. HEAD: bulk load(), torch.load(weights_only=False), cache_values=True.
4.19 Offline MT training loop	assert multitask and task ∈ {mt30, mt80}; load the dataset; for i in range(cfg.steps) (default 10,000,000 iterations): agent.update(buffer), the same update as online (no planning during training). If i % eval_freq == 0 or i % 10,000 == 0: build metrics {iteration, elapsed_time, train metrics}. If i % 50,000 == 0, also run eval (all tasks) and save {i}.pt when i > 0. Log category 'pretrain' with x-axis = iteration; final save at the end. Batch composition: B slices drawn uniformly over all episodes of the pooled dataset, so task frequency follows the dataset's per-task episode counts. There is no per-task balancing; each slice is single-task. Batch size 1024 must be set explicitly (config default 256).	Sec 4.1 p.7 (batch 1024, no other hyperparameter change); App J p.30 (no special offline-RL treatment); Table 1 p.7 (RTX 3090 GPU days: 1M 3.7, 5M 4.2, 19M 5.3, 48M 12, 317M 33)	trainer/offline_trainer.py:67-94; config.yaml:12, 15-16; README (task=mt80 model_size=48 batch_size=1024)	agree (batch composition and the iteration count are code-defined)	PE total_time → HEAD elapsed_time. b67b21c logged at i == 10_000 only (PE: every 10k).
4.20 MT evaluation and normalised score	Offline-trainer eval: for each task_idx, 10 episodes of reset(task_idx), act(obs, t0=(t==0), eval_mode=True, task=task_idx) until done. Logs per-task nanmean episode_reward+<task> and final-step episode_success+<task>, plus the DMC mean (mt30 tasks) and the Meta-World mean reward/success (mt80). Normalised score (evaluate.py): the mean over tasks of 100·success for mw-* and return/10 for DMC, both on [0, 100]. Corrected: evaluate.py never passes eval_mode (default False), so the std[0]·N(0,1) exploration noise is on for every action. It averages with np.mean (not nanmean). Its protocol therefore differs from the trainer eval (eval_mode=True). This holds at HEAD, PE and b67b21c. An mt80 eval is 800 planning episodes.	Sec 4.1 p.7 (normalised score)	trainer/offline_trainer.py:22-40; evaluate.py:69-99 (:80, :91-94); common/logger.py:194-222; tdmpc2.py:203-204	agree on the score definition; the eval protocol is code-defined	—
4.21 ST (online) evaluation	eval_episodes = 10 at step 0 and every 50,000 agent steps, on the training env instance. Per episode: reset; act(obs, t0=(t==0), eval_mode=True) until done. ep_reward = Σ env rewards (each the sum of 2 sim rewards on DMC); success = final-step info['success']. Reported value = np.nanmean over the 10 episodes. eval.csv holds (step, episode_reward) with step = agent step.	Figs 4/5/12/15 (return vs env steps; mean and 95% CI over 3 seeds)	trainer/online_trainer.py:28-52; common/logger.py:224-241; config.yaml:11-12	agree (episode count and frequency are code-defined)	HEAD adds cudagraph_mark_step_begin and an episode_length metric.
4.22 Step units of paper x-axes and results CSVs (resolved)	results/*.csv "step" and the paper's x-axes count simulator (env) steps = agent steps × action repeat. Evidence: MyoSuite CSVs (AR 1) are spaced 50,000 apart, matching eval_freq in agent steps; DMC/MW/MS2 CSVs (AR 2) are spaced 100,000 apart (50k agent steps × 2); dog-run.csv runs to 14,000,000 while the README trains dog-run with steps=7000000 agent steps; footnote 1 p.8: "20k environment steps corresponds to 20 episodes in DMControl". So on DMC "1M" = 500k code (agent) steps = 1000 episodes. A reimplementation that logs agent steps must compare at half the x value. The same applies to Fig 8/17 (20k env steps = 10k agent steps). Eq. 8 calls S "environment steps" although it counts agent steps (a paper inconsistency).	Figs 4, 12, 15; footnote 1 p.8; Eq. 8 p.28	results/tdmpc2/myo-hand-key-turn.csv, walker-run.csv, dog-run.csv; README.md:117; config.yaml:12	agree (once resolved)	—
4.23 Checkpoint contents	agent.save stores only model.state_dict(): world model, π, Q, target/detached Q params, e, masks. Not stored: Adam states (optim, pi_optim), RunningScale S, _prev_mean. A resume or finetune restarts with fresh Adam moments and S = 1.	—	tdmpc2.py:33, 72-79; common/scale.py:11	code_only	Same at PE.
4.24 Custom DMC tasks (needed for mt30/mt80)	11 tasks with custom XMLs (envs/tasks/*.xml, not inspected), registered under 'custom'. walker-walk/run-backwards: stand = (3·tol(torso_h, [1.2, inf), margin 0.6) + (1+upright)/2)/4; move = tol(vx, (−inf, −speed], margin speed/2, value_at_margin 0.5, linear) with speed 1 (walk) / 8 (run); r = stand·(5·move + 1)/6. cheetah-run-backwards: tol(speed, (−inf, −8], margin 8, value 0, linear). cheetah-run-front/back: speed 6; up = (3·foot_up + 2·torso_up)/5 with tol(h, [1.2, inf), margin 0.6); r = up·(5·tol(speed, [6, inf), margin 6, linear) + 1)/6; foot = bfoot (front) / ffoot (back). cheetah-jump: the mean of stand-front and stand-back, each (5·tol((torso_z+foot_z)/2, [1.2, inf), 0.6) + tol(speed, [−0.5, 0.5], margin 0.5, linear))/6. hopper-hop-backwards: tol(height, (0.6, 2)) · tol(speed, (−inf, −1], margin 0.5, value 0.5, linear). reacher-three-easy/hard: 3-link reacher, target size 0.05 / 0.015 (sparse); obs = position, to_target, velocity. cup-spin: not_in_target·(tol(dist, [0.5, inf), 0.25, 0.5, linear) + 2·tol(abs(ball vel), [6, inf), 3, 0.5, linear))/3. Init (corrected): with probability 0.1 any collision-free random position is accepted (ball_x ~ U(−.2,.2), ball_z ~ U(.2,.5)), and it may still be inside the cup. With probability 0.9 the position is resampled until the ball is in the target. So ≥ 90% of episodes start in the cup, and 10% are unconstrained rather than forced outside. pendulum-spin: tol(abs(ang vel), [9, inf), margin 4.5, value 0.5, linear). All tasks: 1000 sim steps (500 agent steps).	App C Table 2 p.18 (11 new tasks with dims)	envs/tasks/walker.py:21-40, 109-128; cheetah.py:21-78, 137-196; hopper.py:33-41, 80-87; reacher.py:25-42, 78-89; ball_in_cup.py:27-99 (:62-70); pendulum.py:21-43; envs/dmcontrol.py:7-10	code_only (the paper gives no reward definitions)	—
4.25 Few-shot finetuning of an MT model	Pretrain 19M on 70 tasks (80 minus 10 held out). Finetune the full model online on one held-out task, starting from an empty buffer with unchanged hyperparameters, for 20k (Fig 8) / 40k (Fig 17) env steps (= 10k / 20k agent steps on DMC). Initialise the new task's e from a source task (Table 7): walker-run ← walker-walk; cheetah-run ← cheetah-run-backwards; hopper-hop ← hopper-stand; pendulum-swingup ← pendulum-spin; reacher-hard ← reacher-easy; bin-picking ← pick-place; box-close ← assembly; door-lock ← door-open; door-unlock ← door-open; hand-insert ← sweep-into. Alternatively use a random vector. Result: normalised score 24.0 from scratch vs 47.0 finetuned (20k steps, 3 seeds). Implementation would need a loader from an MT checkpoint into an ST online loop with that task's mask and e row, plus obs padding to the pretraining max dim.	Sec 3.3 p.5; Sec 4.1 p.7-8, Fig 8, footnote 1; App E p.24, Fig 17, Table 7	not implemented (no 70-task set; grep 'finetun' finds nothing; train.py:52; evaluate.py:49-51)	paper_only	—
4.26 Test-time regularisation	§3.26.	App J p.30	—	paper_only on main	—
4.27 DMC ST sanity targets (released curves)	From results/tdmpc2/<task>.csv (columns step, reward, seed; points every 100k env steps). Mean over seeds at env steps 100k/500k/1M/2M/4M (Dog/Humanoid also at 14M): acrobot-swingup 179/361/518/611/663; cartpole-balance 998/994/975/999/998; cartpole-balance-sparse 1000/1000/1000/1000/977; cartpole-swingup 839/876/876/876/877; cartpole-swingup-sparse 0/845/842/849/850 (n=2 at 4M); cheetah-jump 576/709/803/809/823; cheetah-run 519/758/844/860/896; cheetah-run-back 462/711/800/833/846; cheetah-run-backwards 520/680/706/724/754; cheetah-run-front 311/666/677/675/694; cup-catch 981/985/986/984/984; cup-spin 843/848/847/848/848; dog-run 27/133/170/427/682/870; dog-stand 147/700/799/933/949/980; dog-trot 12/132/500/884/947/951; dog-walk 35/483/494/885/943/976 (Dog 14M values n=2); finger-spin 948/985/986/989/991; finger-turn-easy 488/855/918/981/982; finger-turn-hard 484/876/975/918/977; fish-swim 205/610/690/706/761; hopper-hop 14/303/338/375/449; hopper-hop-backwards 273/891/902/913/932; hopper-stand 28/936/954/954/955; humanoid-run 1/108/185/316/461/603; humanoid-stand 6/422/664/913/918/906; humanoid-walk 1/278/628/860/912/900; pendulum-spin 838/845/846/845/843; pendulum-swingup 724/842/819/837/838; quadruped-run 239/940/938/948/954; quadruped-walk 487/957/956/970/971; reacher-easy 937/919/983/982/987; reacher-hard 414/914/981/980/982; reacher-three-easy 880/939/944/970/982; reacher-three-hard 267/681/720/874/846; walker-run 723/820/842/862/877; walker-run-backwards 516/795/866/904/939; walker-stand 981/957/991/992/993; walker-walk 962/979/980/984/981; walker-walk-backwards 953/978/977/984/984. 39-task mean: 482.5 @100k, 736.2 @500k, 792.8 @1M, 852.7 @2M, 878.1 @4M env steps. Seed spread at 1M (min/med/max): walker-run 834/837/856; cheetah-run 725/897/911; quadruped-run 926/941/947; hopper-hop 271/369/373; acrobot 477/484/594; humanoid-walk 610/632/642. Baselines at 1M env steps (walker-run / cheetah-run / quadruped-run / hopper-hop / acrobot / humanoid-walk): TD-MPC 824/830/894/319/380/≈660 (corrected: humanoid-walk has 5 rows at 1M with duplicated seed labels: 636.2, 706.9 (s1), 745.3, 686.1 (s2), 526.5 (s3). The mean of all rows is 660.2 and per-seed means give 638, so the earlier "640" is not reproducible.) DreamerV3 581/638/768/200/327/1 (humanoid-walk: a single row); SAC 635/730/257/70/10/2.	App D Fig 12 p.21 (DMC, first 4M), Fig 15 p.23 (locomotion to 14M); Fig 4 p.5; Table 2 p.18	results/tdmpc2/*.csv (added in b67b21c, renamed in 88095e7); results/{tdmpc,dreamerv3,sac}/*.csv	agree	Produced with b67b21c-era code and physics (mujoco 2.3.1). HEAD may deviate slightly (mujoco 3.1.2, TF32, entropy form).
4.28 MT sanity targets	Normalised score (0-100) after offline training. mt80: 1M 16.0, 5M 49.5, 19M 57.1, 48M 68.0, 317M 70.6 (Table 1; GPU days 3.7/4.2/5.3/12/33). mt30: 1M 18.9, 5M 28.3, 19M 54.2, 48M 59.4, 317M 71.4 (Fig 7 / Fig 22). 15-task DMC subset: 61.7, 78.0, 80.5 (three sizes at about 5M/19M/48M; the task list is not given). 19M mt80 ablations (Fig 9 p.8): actor: policy 42.2, planning 53.7, planning+policy 54.2; normalisation: none 46.8, SimNorm 51.0, LN+SimNorm 54.2; regression: continuous 49.6, discrete 54.2; #Q: 2 → 53.5, 5 → 54.2, 10 → 57.0. Task-embedding normalisation (Fig 18): 46.6 unnormalised vs 54.2. Few-shot: 24.0 → 47.0. The 19M mt80 score is 57.1 (Table 1), 54.2 (ablations) and 56.54 (Table 10), so expect about ±3 run-to-run variation.	Table 1 p.7; Figs 7, 8 p.7; Fig 9 p.8; App F Fig 18 p.25; App K Fig 22 p.31	evaluate.py:94	agree	—
4.29 Hyperparameter cross-check (Table 8 vs config)	All values agree except the encoder lr (1e-4 vs 9e-5). Table 8 "Policy prior loss norm.: Moving (5%, 95%) percentiles" is RunningScale (§2.13), whose EMA rate τ = 0.01 is code_only. Code-only defaults: eval_freq 50,000; eval_episodes 10; steps 10M; termination_coef 1 (HEAD); compile true (non-algorithmic). The full table is in §5.	Table 8 p.27	config.yaml:1-91; tdmpc2.py:22-31	conflict (encoder lr only)	config.yaml at HEAD adds episodic, termination_coef and compile; disable_wandb became enable_wandb. Hyperparameters are unchanged vs PE.
4.30 Model size per experiment	ST online, all 104 tasks: 5M. mt80 and mt30 scaling: all 5 sizes. Ablations (Fig 9) and Table 10: 19M on mt80. Few-shot: 19M on 70 tasks. 15-task DMC: 3 sizes.	Sec 4 p.6 (5M unless stated); Fig 9 caption p.8; App E p.24; App J p.30; App K p.31	common/__init__.py:1-24	agree (task_dim/latent quirks in §1.20-1.21)	—
4.31 Numerics and dtypes (loop side)	All data is f32 (obs, action, reward, terminated). episode id is int64; task id is .long() before lookup and used as an integer index into the discount and mask tables.	not specified	envs/wrappers/tensor.py:19-42; common/world_model.py:96	code_only	—

4.C Reference pseudocode: online loop (HEAD)

step = 0; done = True; eval_next = False                       # [PE: eval_next = True]
while step <= steps:                                            # steps+1 iterations (default 10_000_000)
    if step % eval_freq == 0: eval_next = True                  # eval_freq = 50_000
    if done:
        if eval_next: eval(); eval_next = False                 # at the NEXT episode boundary; includes step 0
        if step > 0:
            if info['terminated'] and not episodic: raise ValueError
            log(episode_reward = sum(r over tds[1:]), episode_success = info['success'] (last step))
            buffer.add(cat(tds))                                # whole finished episode, T+1 rows
        obs = env.reset(); tds = [row(obs, NaN action, NaN reward, NaN terminated)]
    action = agent.act(obs, t0 = (len(tds) == 1)) if step > seed_steps else env.rand_act()
    obs, reward, done, info = env.step(action)
    tds.append(row(obs, action, reward, info['terminated']))
    if step >= seed_steps:
        n = seed_steps if step == seed_steps else 1             # one-off pretraining burst
        for _ in range(n): agent.update(buffer)
    step += 1
save()                                                          # in-progress episode is discarded

4.A Implementation subtleties (training loop, replay, envs, MT)

• Slices never cross episodes. Row 0 placeholders are never consumed (losses use rows 1..H). In JAX, store zeros, not NaN.
• Insert data per completed episode, or sample only fully written windows. The burst at step == S must see exactly the completed seed episodes (5 × 500 transitions on DMC).
• Random actions run for steps 0..S inclusive. Updates start at step == S with S back-to-back updates, then 1 per agent step.
• Use an episode ring buffer [N_eps, T+1, ·] with N_eps = 1e6 // T (2000 for DMC) plus randint(episode) and randint(start ∈ [0, T−H]). This is the original b67b21c design and is jit-friendly.
• The time limit is not terminal. Bootstrap from the final stored obs. terminated is 0 everywhere at PE semantics.
• DMC: action_scale to [−1, 1], 2 sim steps per agent step with summed reward, flattened obs dict in key order as f32, no obs/reward normalisation, task_kwargs={'random': seed}.
• Eval: 10 episodes at step 0 and every 50k agent steps, with eval_mode=True. eval_mode removes only the final noise.
• Compare against the paper and results CSVs at env steps (= 2 × agent steps on DMC).
• Represent MT machinery as data (a mask table, a discount table and an embedding table indexed by task id), not as agent kwargs, which fits the Ajax Extension rule.
• MT: offline only, batch 1024, uniform over pooled episodes, no per-task balancing, no planning during training.
• Neither Adam state nor RunningScale is checkpointed in the reference. Decide explicitly whether Ajax checkpoints them (better for resume, but differs from the reference).

4.B Open questions (training loop, replay, envs, MT)

• UTD with vectorised GPU envs: the reference is single-env with one update per agent step. Use the paper-consistent mapping of n_envs updates per vector step (one per collected transition), with episode-wise insertion?
• Offline iteration count for the paper's MT models: the paper gives only GPU days; the code default is 10M iterations at batch 1024.
• Dataset internals (key names, dtypes, task id dtype, zero-padding of invalid action dims, how DMC 500-step episodes became 101-row segments in mt80): these cannot be inspected because the data is external.
• torchrl SliceSampler at the pinned versions (strict_length raise vs drop, circular wrap): unverified. It does not matter for fixed-length episodes.
• Which version's quirks to reproduce (policy-gradient leak into e, entropy form, PE clip quirk): the recommendation is PE semantics minus the clip quirk (user decision).
• Meta-World at HEAD appears broken (read-only max_episode_steps). The reset zero-action quirk matters only for mt80.
• The 15-task DMC subset (App K) is not listed anywhere.
• Seed-phase action RNG: the reference does not seed it. Ajax will seed it via JAX keys. Confirm that this deviation is fine.
• Few-shot finetuning: in scope? There is no reference code.

---

5. Hyperparameter defaults (flat table)

Scope: ST = single-task online (5M, DMC defaults); MT = multi-task offline (mt30/mt80). The last column flags every HEAD (e9f5932) vs PE (5f6fade) difference and says which one the paper describes. "=" means identical at HEAD and PE.

Name	Value	Scope	Source	HEAD vs PE; what the paper describes
model_size	5 (M)	ST	Sec 4 p.6; README	= (b67b21c default 5; PE/HEAD ???, must be passed)
model_size	1 / 5 / 19 / 48 / 317	MT	Table 9 p.28	=
enc_dim	256 (1M 256, 19M 1024, 48M 1792, 317M 4096)	ST/MT	Table 8/9; common/__init__.py	=
mlp_dim	512 (1M 384, 19M 1024, 48M 1792, 317M 4096)	ST/MT	Table 8/9	=
latent_dim	512 (1M 128, 19M 768, 48M 768, 317M 1376; mt30-19M: 512 in code)	ST/MT	Table 8/9; parser.py:67-68	=; paper says 768 for 19M (conflict, code quirk at both)
num_enc_layers	2 (19M 3, 48M 4, 317M 5)	ST/MT	Table 9	=
num_q	5 (1M 2, 317M 8)	ST/MT	Table 8/9	=
Q subsample size	2 (always)	all	Sec 3.1; Table 9 caption	=
dropout (Q, first layer only)	0.01	all	Table 8; App A	=
simnorm_dim V	8	all	Table 8	=
SimNorm temperature τ	1 (implicit; not in code)	all	Table 8	=
task_dim	0	ST	parser.py	=
task_dim	96 (paper, all sizes); code: 96 for mt80 or sizes 1/317, 64 for mt30 at 5/19/48M	MT	Table 8/9; parser.py:75	=; paper describes 96 (conflict at both pins)
task-embedding max norm	1	MT	Table 8	=
task-embedding init	U(−0.02, 0.02)	MT	init.py:10-11	= (code_only)
embedding renorm eps	1e-7	MT	torch max_norm internals	=
LayerNorm eps	1e-5	all	torch default	= (code_only)
Linear weight init std	0.02 (trunc bounds ±2 absolute, effectively untruncated)	all	init.py:7	= (HEAD restored in e9f5932; regression window 2024-09-25 to 2026-07-13)
Linear bias init	0	all	init.py	=
Zero-init	reward final weight; every Q member's final weight	all	world_model.py:30-32	=
num_bins K	101	all	Table 8	=
vmin, vmax (symlog space)	−10, +10	all	config.yaml:51-52	= (code_only)
bin_size	0.2	all	derived	=
log_std_min, log_std_max	−10, 2	all	Table 8	=
entropy_coef β	1e-4 (effective ×#valid action dims)	all	Table 8	value =; form differs: PE keeps the tanh-Jacobian (unscaled) and subtracts ½ln2π once; HEAD drops the tanh gradient and subtracts ½ln2π per dim. The paper describes neither (closed form, no scaling).
Gaussian log-prob constant	0.9189385175704956 (½ln2π)	all	math.py	HEAD: per dim (incl. masked); PE: once, then ×n
squash eps	1e-6	all	math.py:27	=
entropy_scale eps	1e-8	all	world_model.py:176	HEAD only (PE has no such term); paper n/a
horizon H	3 (planning and training rollout)	all	Table 8	=
rho (λ)	0.5	all	Table 8	= (HEAD mujoco-* override 0.7, HEAD-only, n/a to the paper)
consistency_coef	20	all	Table 8	=
reward_coef	0.1	all	Table 8	=
value_coef	0.1	all	Table 8	=
termination_coef	1	episodic only	config.yaml:19	HEAD only (PE: absent); paper: absent
episodic	false	all	config.yaml:7	HEAD only flag; paper: no terminations (= PE)
lr	3e-4	all	Table 8	=
enc_lr_scale	0.3 → encoder lr 9e-5	all	config.yaml:23	=; paper says encoder lr 1e-4 (conflict at both pins)
Adam betas	(0.9, 0.999)	all	torch default	= (paper unspecified)
Adam eps (world model)	1e-8	all	torch default	= (paper unspecified)
Adam eps (policy)	1e-5	all	tdmpc2.py:31	= (paper unspecified)
weight decay	0	all	torch default	=
grad_clip_norm	20 (world model global; policy separate)	all	Table 8	value =; PE world-model norm also includes stale policy grads, HEAD does not; paper unspecified
clip eps	1e-6	all	torch	=
tau	0.01 (target-Q EMA = 1 − 0.99; also RunningScale rate)	all	Table 8 (momentum 0.99); config.yaml:25	=
RunningScale init	1.0	all	scale.py:11	= (code_only)
RunningScale percentiles	5, 95	all	Table 8	=
RunningScale min range	1	all	scale.py	= (code_only)
batch_size	256	ST	Table 8	=
batch_size	1024 (set via CLI)	MT	Sec 4.1 p.7; README	=
UTD	1	ST	Table 8	=
discount_denom	5	all	Eq. 7	=
discount_min	0.95	all	Eq. 7	=
discount_max	0.995	all	Eq. 7	= (HEAD mujoco-* override 0.99, HEAD only)
γ (DMC, T=500)	0.99	ST/MT	Eq. 7	=
γ (MW/MS2/Myo, T=100)	0.95	ST/MT	Eq. 7	=
mpc	true	all	config.yaml:34	= (b67b21c had no flag; always planned)
iterations	6 (+2 if A ≥ 20 → 8; MT uses padded A = 6 → 6)	all	Table 8	=
num_samples N	512 (includes the 24 π trajectories)	all	Table 8	=
num_pi_trajs P	24	all	Table 8	=
num_elites E	64	all	Table 8	=
min_std	0.05	all	Table 8	=
max_std	2 (also the per-step initial std)	all	Table 8	=
temperature	0.5 (multiplies: exp(0.5·ΔV))	all	Table 8	=
MPPI momentum	none	all	Table 8; App A	=
MPPI denominator eps	1e-9	all	tdmpc2.py:192-193	=
action bounds	[−1, 1]	all	action_scale; tdmpc2.py:178, 206	=
buffer_size	1,000,000; ST capacity = min(buffer_size, steps)	ST	Table 8	Units differ: b67b21c = 1e6//T episodes (2000 DMC episodes = 1M transitions; produced the paper's curves); PE/HEAD = 1e6 rows (≈998k transitions). The paper says "1,000,000" (b67b21c matches exactly).
buffer capacity	345,690,000 rows (mt30); 550,450,000 rows (mt80)	MT	offline_trainer.py:54	= from PE, but the PE assert fails (rows vs episodes); b67b21c counted episodes (690,000 / 5,450,000); HEAD warns. Paper: 345M / 545M transitions.
dataset episode rows L	501 (mt30), 101 (mt80)	MT	offline_trainer.py:53	=
slice length	H+1 = 4 rows	all	buffer.py	= (b67b21c: RandomCrop(H+1))
seed_steps S	max(1000, 5T) → 2500 (DMC), 1000 (T=100)	ST	Eq. 8	=
pretraining burst	S updates at step == S	ST	online_trainer.py:115-123	= (code_only)
steps	10,000,000 (agent steps online; iterations offline)	ST/MT	config.yaml:15	= (code_only; the paper reports 4M-14M env steps per domain)
eval_freq	50,000 (agent steps / iterations)	ST/MT	config.yaml:12	= (code_only)
eval_episodes	10 (per task in MT)	ST/MT	config.yaml:11	= (code_only)
offline log interval	10,000 iterations	MT	offline_trainer.py:76	PE/HEAD every 10k; b67b21c only at i == 10,000
seed	1	all	config.yaml	=
DMC action repeat	2	ST/MT	App C Table 6	= (wrapper implementation differs; equivalent)
DMC episode length	500 agent steps (1000 sim)	ST/MT	App C Table 6	=
MW / MS2 / Myo episode length (action repeat)	100 (2) / 100 (2) / 100 (1)	ST/MT	App C Table 6	ManiSkill2: HEAD may terminate early; PE fixed length (paper = PE)
MT padded dims	mt30: obs 24, act 6; mt80: obs 39, act 6	MT	common/__init__.py; wrapper	=
float32 matmul precision	'high' (TF32)	all	train.py:22-23	HEAD only; PE 'highest' (full fp32); paper unspecified
compile	true (performance only)	all	config.yaml	HEAD only; no math change
obs	'state'	ST/MT	config.yaml:6	= (b67b21c state only)
num_channels	32 (RGB only; out of scope)	—	config.yaml	= (PE added it)