| D1 |
DreamerV3 |
Asynchronous: acting parameters ≈ 2 driver iterations stale; one-batch prefetch; write-back and metrics one train call late; first batch prefetched before the gate |
Synchronous |
Infrastructure, not algorithmic |
| D2 |
DreamerV3 |
Replay sampler RNG fixed to seed 0; the report stream pops the online queue every 180 s of wall-clock, and its loss pass reads the return normaliser without updating it (agent.py:410) |
Seeded from the run seed; no report stream, so no non-updating loss pass |
Infrastructure |
| D3 |
DreamerV3 |
Updates interleaved between the per-env adds of one vector step |
Updates after the tick's adds |
Same rate; intra-tick order only |
| D4 |
DreamerV3 |
bf16 compute (2411f7d: bf16 RMS statistics, f32 slow critic) |
f32 |
Numerics only; bf16 revisited at the GPU milestone |
| D5 |
DreamerV3 |
One decoder head per observation key (dict observations) |
One head for the flat observation vector |
Loss identical (sum of means); the per-tensor AGC partition differs (≤ 8 % clipped-gradient difference, probed) |
| D6 |
DreamerV3 |
Optimizer and normaliser state checkpointed; Ratio state, online queue and replay RNG not |
Everything checkpointed |
Resume fidelity |
| T7 |
TD-MPC2 |
Single environment |
n_envs > 1: per-env t0 and warm-start mean, n_envs updates per stepping tick, seed phase / burst rule of DESIGN.md §4.4 |
Generalisation; exact for n_envs = 1 |
| T8 |
TD-MPC2 |
The first planned step after seeding warm-starts from the step-0 evaluation's leftover mean |
Zeros |
Negligible |
| T9 |
TD-MPC2 |
Evaluation reuses the training env instance and planner mean; the first evaluation is at step 0 (untrained agent) |
Rebuilt evaluation env, separate planner carry; each evaluation from fresh initial states (its key folded with the number of evaluations so far: deterministic per seed, fresh across evaluations as the reference's resets of its running env are); evaluations every log_frequency env steps from the start of training, the first after log_frequency steps (on the held tick ending an episode when log_frequency is a multiple of T·n_envs, as the reference's eval_freq is) |
Evaluation protocol only |
| T10 |
TD-MPC2 |
Paper era: no termination guard (any done treated as a time limit); latest: ValueError unless episodic |
Terminations counted (episode ends off the T-step schedule, and true terminations on its last step); ValueError raised after the run; episodic=True takes them instead (T29) |
Loud failure instead of silently wrong targets |
| T11 |
TD-MPC2 |
Seed-phase actions unseeded (gym space RNG) |
Seeded |
Reproducibility |
| T12 |
TD-MPC2 |
Adam state and RunningScale not checkpointed |
Checkpointed |
Resume fidelity |
| T13 |
TD-MPC2 |
Matmul precision: TF32 (latest) / full fp32 (paper era) |
highest in the parity tests; JAX's default elsewhere (not yet set for GPU runs) |
Numerics |
| T14 |
TD-MPC2 MT |
task_dim 64 for mt30 at 5/19/48M (an acknowledged accident) |
96 (paper) |
Ajax's task set is not mt30 |
| T15 |
TD-MPC2 MT |
Evaluation-time embedding renorm persisted in place |
Renorm applied but not persisted at evaluation |
Evaluation cannot write parameters |
| T16 |
TD-MPC2 MT |
evaluate.py normalised score with exploration noise on; the offline trainer evaluates after its first update, then every eval_freq |
Offline-trainer protocol (eval_mode, 10 episodes per task, t0 at each episode's first step), every log_frequency updates of the run, the first after log_frequency updates; each evaluation from fresh initial states; normalised score mean(return / 10) (the DMC normalisation) |
The trainer protocol is the one behind the paper's training curves |
| T17 |
TD-MPC2 MT |
Datasets mt30 / mt80 (replay of 240 single-task agents, random to expert) |
Self-generated: full training history of single-task Ajax runs on the playground versions of the 19 original mt30 DMC tasks |
Validates the mechanisms, not the paper's numbers |
| E18 |
both |
dm_control (MuJoCo C) |
mujoco_playground MJX ports: reduced solver iterations, some sim_dt changes (e.g. reacher 0.005 vs 0.02), impl='jax' on CPU. QuadrupedRun / DogRun and the 11 custom TD-MPC2 tasks are unavailable |
Environment, not algorithm |
| E19 |
both |
dm_control action repeat stops at the episode end |
brax EpisodeWrapper repeat continues after a termination |
None on DMC (never terminates) |
| E20 |
both |
Fresh random initial state every episode |
Same: static reset mode resets freshly itself; dynamic mode uses the auto-reset, fresh on every backend since playground envs default to fresh_reset=True (#54). Only a playground env built with fresh_reset=False restarts from a cached first state in dynamic mode (warned) |
None with the default env stack |
| E21 |
both |
— |
Golden checksums are CPU-only regression guards |
GPU nondeterminism (TF32, scatter-add order) |
| D22 |
DreamerV3 |
Two-hot bins symexp(jnp.linspace(-20, 0, 128)) evaluated in float32 inside the jitted step; expectation p_m·b_m + Σ (p_i·b_i + p_j·b_j) over mirror pairs, which XLA fuses into multiply-adds under jit, so a zero-initialised head predicts a constant that depends on the compiled program (−0.17 to +0.16 on CPU), not the 0 of paper p.18 |
Bins: float32 rounding of the float64 values, computed once with numpy (the same constants eagerly, under jit and on every backend); expectation Σ (p_j − p_i)·b_j over mirror pairs, exactly 0 for uniform p also under jit |
Same expectation in exact arithmetic: bins within 1.3e-6 relative of the reference's; encode and loss bit-identical given the same bins; decode within 3·ε·E_p|b| (ε = float32 epsilon). Not negligible in training. Once the logits move, the reference's contracted pair sum adds noise of std ≈ 0.085 (max ≈ 0.3) to every reward, value and slow-critic prediction, plus a constant at zero logits that depends on the compiled program (−0.17 to +0.16); the noise fades as the outer bins lose mass (CartPole, 1m: value std 0.085, 0.052, ≈ 0.01 over the 2000-, 4000- and 6000-row windows). It slows the early decline of the actor's entropy (CartPole, 1m, 4000-row window: 0.51 with the reference's sum, 0.46 with Ajax's, measured in both codebases); with the same expectation on both sides (and the replay sampler seeded per run), the training statistics' seed ranges overlap the reference's in at least 52 of the 58 windows with updates (PERFORMANCE_REPORT.md, DreamerV3 reference comparison, round 3, for the exceptions). Ajax's form is exact only while the outer logits stay bitwise mirror-symmetric, which depends on the compiled program (a one-seed run and a three-seed vmapped run differ): a float32 fragility of the bins (±4.85e8, holding ≈ 1/255 of the mass early on) that both codebases share |
| D23 |
DreamerV3 |
Kernel init 1.1368 · sqrt(1/fan_in) · TruncatedNormal(−2, 2) |
variance_scaling(outscale², 'fan_in', 'truncated_normal'), i.e. the exact factor 1/0.87962566 |
Same draws up to a constant factor: 4.16e-5 relative |
| T24 |
TD-MPC2 |
Two-hot over torch.linspace(vmin, vmax, 101) (middle bin −1.5e-7, not 0), any vmin < vmax; target by the floor formula k = floor((x − vmin)/Δ); naive expectation Σ p_k·b_k; F.log_softmax |
Bins: float32 rounding of the float64 grid, mirrored (within 4.8e-7 of torch's); target by linear interpolation between the neighbouring bins; symmetric expectation as in D22; logits − logsumexp(logits); symmetric ranges only (vmax = −vmin, as in every TD-MPC2 config) |
Float32 rounding: weights within 1.1e-5 of the floor formula and closer to the exact ones (2e-6 vs 8.6e-6); decode within ε·(1 + 6·E_p|b|) in symlog space (≤ 4.8e-6); log-probabilities within 3·ε·(max|logits| + loss) |
| T25 |
TD-MPC2 |
RunningScale: sort-and-interpolate _percentile; S.lerp_(v, τ), i.e. S + τ·(v − S) |
jnp.percentile (linear: the same formula); (1 − τ)·S + τ·v (optax.incremental_update) |
Float32 rounding (≤ 1e-6 relative over 30 updates, tested) |
| T26 |
TD-MPC2 |
b67b21c replay capacity floor(min(buffer_size, steps) / T) episodes (buffer.py:40) |
ceil(min(buffer_size, total) / (T·n_envs)) rounds of n_envs episodes, plus one staging round that is written in place and never sampled; total = the run's env steps when the current train() call ends, so a resumed run (chunked training, or a checkpoint restored by a new agent) sizes its ring for the whole run and moves the carried episodes into it (EpisodeBuffer.adopt) |
Identical when buffer_size is a multiple of T·n_envs (DMC: 1e6 / 500 = 2000 episodes) and whenever the run never fills the buffer; otherwise at most one more round. A run trained in chunks replays exactly what the uninterrupted run does |
| T28 |
TD-MPC2 |
One update per env step after the seed-phase burst (online_trainer.py:105-112); config.yaml has no update-to-data option |
updates_per_step r (default 1): max(1, round(n_envs·r)) updates per stepping tick after the burst, which stays seed_steps updates; in the episodic mode (T29) max(1, round(m·r)) on a tick that stepped m > 0 envs |
Opt-in generalisation; r = 1 is the reference exactly (and T7 for n_envs > 1) |
| T29 |
TD-MPC2 |
Latest code's episodic mode (eece801, as at e9f5932): single env; finished episodes added to a row storage of min(buffer_size, steps) rows; SliceSampler(strict_length=True): uniform episode among those of ≥ H+1 rows, uniform offset; seed phase and one update per env step counted in env steps; every other part of the agent at its latest version |
Opt-in (episodic=True), grafted onto the paper-era agent (every §2 choice above kept, e.g. Q dropout in the TD target); dynamic-reset row collector: per-env rings of ceil(min(buffer_size, total) / n_envs) rows written in lockstep (the running episode's rows included; total = the run's rows when the current train() call ends, the reference's clamp to the run, so a resumed run moves the carried rows into its own rings, EpisodicBuffer.adopt), an overwritten episode keeps its stored tail as in the reference; seed phase max(S div n, T) ticks (held rows included, ≈ S·k/(k+1) env steps for k-step episodes), then one update per env stepped in the tick (none until an episode of ≥ H+1 rows is stored); n_timesteps counts rows |
Same targets, losses, planning and sampling distribution; the seed phase is shorter by the episode ends it contains. Validated on a terminating chain whose values are known (test_tdmpc2_episodic.py) |
| D26 |
DreamerV3 |
when.Ratio keeps prev in float64: prev += repeats / r |
The exact closed form 1 + ⌊(t − t0)·r⌋ on the absolute tick (r the exact fraction of train_ratio / (B·T)) |
Identical whenever 1/r is exact in float64 (every reference configuration: r = 2^k); otherwise the reference's rounding sometimes runs an update one transition late, and Ajax is at most one update ahead (tested) |
| D27 |
DreamerV3 |
Replay capacity in items (5e6), FIFO over the items of all workers |
Per-env rings of C = min(ceil(replay_capacity / n_envs), rows per env of the run) rows, overwritten in lockstep |
Never binds in the paper protocols (DMC: 250K rows ≪ 5e6), in one run or split into resumed calls (a resumed run grows a run-clamped ring to its total length); when it binds, every env drops its oldest row together and L − 1 fewer items per env are kept than rows |