Status: trained. 24,797,184 simulated steps, 90,000 updates. Weights, optimizer state,
configs and logs are in this repo.
What it is
A planning-oriented latent world model that predicts future latent states, rewards and values from raw
(non-privileged) observations, so that a sampling-based planner can rank candidate action sequences.
z_t = E(o_t) latent encoder, newest observation only
h_t = C(o_{t-H:t-1}, a_{t-H:t-1}) history encoder — proprio + actions only, no images
ẑ_{t+1:t+T} = f(z_t, a_{t:t+T-1}, h_t) causal transformer, whole horizon in one pass
r̂_{t:t+T-1} = R(ẑ, a) sequence-to-sequence transformer head
v̂_{t+1:t+T} = V(ẑ) sequence-to-sequence transformer head
â_{t:t+H} = π(z_t, h_t) base policy, action chunks to warm-start planning
Keeping images out of the history encoder is the paper's Minimal History Representation; it is what
makes planning affordable. Predicting the horizon in a single forward pass, rather than unrolling, is
what makes sampling thousands of candidate trajectories tractable.
Architecture
Paper Table II. Embedding dimension 64; all transformer MLPs hidden 256; dynamics 3 layers / 4 heads;
reward 1 / 1; value 1 / 1; base policy 4 layers / 8 heads. Horizons H = T = 5.
The encoder passes each of three camera views through an ImageNet-pretrained ResNet-18 (shared trunk by
default, 11.18 M parameters against 33.53 M for separate trunks) to 3×512, concatenates with the 20-d
proprioceptive observation, and projects to the 64-d latent. The torchvision→Flax NNX weight conversion
is exact: max absolute deviation 1.9e-06.
Training objective
Four terms, weighted 1 / 1 / 1 / 4:
latent dynamics — MSE against stop_grad(E(o_{t+i+1}))
reward — MSE
value — MSE against the expert critic's V
behaviour cloning — MSE, masked by the cumulative expert flag so it stops contributing the
moment an environment leaves the expert
There is deliberately no pixel-reconstruction loss. The paper's ablation shows adding one drops
manipulation success from 0.90 to 0.32 — reconstruction pressures the latent to encode randomised
texture and lighting that are irrelevant to the task.
Read this before using the numbers
No success rate is reported here, and none of these numbers is comparable to the
paper's Table I. Table I reports task success, which requires closed-loop rollouts
under the MPPI planner. The planner exists in the port (simdist/control/mppi.py, task
agnostic) but the manipulation closed-loop harness does not — only the locomotion one
(scripts/simulate_go2.py). Everything below is a training-side proxy: latent
dynamics error, head regression quality, and a value-ranking AUC.
Results at step 90,000
metric
value
reference
test/latent_dynamics
0.0278
paper reports 0.076 / 0.019
train/latent_dynamics
0.0343
eval/latent_rollout_error h1→h5
0.0164, 0.0228, 0.0283, 0.0344, 0.0385
mean 0.0281, no knee
eval/value_pearson_r
0.9510
data critic ceiling +0.953
eval/value_r2
0.9040
eval/reward_pearson_r
0.9300
eval/reward_r2
0.8641
eval/bc_action_mse
0.1188
eval/value_auc_success
0.7344
see caveat below
eval/latent_variance encoder / predicted
0.3695 / 0.3482
matched — no collapse, no runaway
The value head is the informative one. It regresses the data-generating critic's output,
so that critic's own correlation (r = +0.953 against bootstrapped return-to-go) is a
ceiling it cannot meaningfully exceed. At 0.951 it has essentially reached it.
The AUC is optimistic.eval.episodes.holdout is 0, so the train/test split is at the
chunk level and no episode is fully unseen — the AUC has a step-level leak.
eval/value_auc_holdout_episodewise: 0 records that regime. It is also computed over only
32 episodes.
Stability across training
Early rows were transcribed from live reads during the run; auto-park pulled only the log
tail, so the shipped logs/train_metrics.txt starts at step 77500.
step
latent_dynamics
encoder variance
lr
value AUC
2500
0.424
0.309
5.0e-5
0.766
5000
0.080
0.243
1.0e-4
0.684
7500
0.084
0.245
1.5e-4
0.731
10000
0.076
0.336
2.0e-4 (peak)
0.773
90000
0.034
0.369
1.0e-4
0.734
Encoder variance staying flat through and past peak LR is the load-bearing observation —
see the divergence section.
The divergence, and the deviation it forced
The first attempt at this run diverged and had to be thrown away.latent_dynamics
went 9.37 → 8.62 → 3.0e20 over steps 2500/5000/7500 and settled near 1e28, while reward,
value and action losses stayed bounded at 0.5–2.7 the whole time. latent_variance/encoder
tracked it exactly: 0.97 → 2.88 → 3.1e13 → 1.6e25.
The cause is a degenerate direction in the latent objective. The dynamics target is
stop_gradient(encode_latent(...)), but it is produced by the same online encoder — an
encoder that grows also grows its own target. stop_gradient bounds the gradient path, not
the magnitude. Nothing in the loss penalises ‖z‖, and the encoder's latent_mlp ends in a
free linear map, so ‖z‖ is free to run away. The optimizer was bare optax.adam, with
global_norm computed for logging and never applied, so nothing bounded the rate either.
Two fixes were probed at peak LR 2e-4 held for 4000 steps (warmup compressed to 1000):
configuration
latent_dynamics
encoder variance
value AUC
gradient clipping only
5.4 → 9.4e6
0.99 → 1.1e4
0.80 → 0.57
clipping + latent LayerNorm
0.63 → 0.063
pinned 0.12–0.45
~0.70–0.80
Clipping alone only bounds how fast the run travels the degenerate direction. The LayerNorm
removes it. This model therefore uses a LayerNorm on the encoder latent, which the paper
does not describe — an explicit deviation, exposed as model.encoder.latent_norm
(default false, the paper's architecture) rather than hardcoded. Whether this is an
undocumented detail of the reference implementation or a mis-port has not been checked
against the reference code.
Deviations from the paper
Latent LayerNorm — added, as above. The substantive one.
Gradient clipping at global norm 1.0. The paper logs grad_norm and does not apply it.
90,000 updates, roughly 1.06 epochs, against the paper's ~194.5 k / 2 epochs. A budget
decision: per Table I, data scale dominates and epochs are secondary, so the full dataset
with fewer passes is the better trade under a cap.
H = T = 5, not the paper's 25. Deliberate and documented in the port: this policy runs
at 10 Hz, so 5 steps is 0.5 s of history and prediction — the timescale an insertion
evolves on. 25 would be 2.5 s.
No pixel-reconstruction loss and no image decoder, matching the paper (Table I: adding
reconstruction as an objective is 0.90 → 0.32).
Adam, cosine 2e-4 → 1e-4 with 10 k warmup. Loss weights 1/1/1/4 (latent dynamics / reward /
value / behaviour cloning, the last masked by the cumulative expert flag). One RTX 5090,
2.26 updates/s, 11.9 h, GPU util 97 %, dataloader_wait_frac 0.015.
Contents: checkpoint/ is an orbax checkpoint including optimizer state (resumable);
config/ holds the model, system and training configs; logs/ and metrics.json hold the
eval history that survived.
Data
24,797,184 steps, 155,370 episodes, 843 shards merged from 116 independent generation runs
(37.2 h on one RTX 5090, 174 rows/s). Three RGB views at 160×120 (front, side, wrist), JPEG
q90, plus 20-d proprioception and 7-d relative Cartesian OSC actions. expert_prob 0.5,
sub-optimal actions drawn from a 37-checkpoint ladder, visual and physics domain
randomisation on.
The dataset is not published. It lived on rented storage that was released at the end of
the campaign; only the model and its logs were retained.
Known data defects
Both are recorded because they affect anyone reproducing this, and neither caused the
divergence above (reward and value losses stayed bounded throughout):
Value outliers. 0.116 % of rows have value < −20, p0.01 = −94.97, against mean 8.87
and std 4.00 — a −28σ target after scaling. This matches an MDP-mismatch signature seen
earlier in the project (V^e −98 vs +8.95).
Action outliers.ee_delta_* are expected to be order ~1, but per-dimension absolute
maxima reach 3788 while p99.9 is only 12–37, inflating actions.std to 7.2–13.7.
A further note for reproducers: episode_ids from the generator is a per-environment
counter that restarts at 0 in every process, so merging N generation runs puts N distinct
episodes under (env 0, episode 0). They must be namespaced per source run before
processing, or the episode stitcher will refuse the merged set.
Intended use
Sim-only evaluation: held-out prediction losses, latent rollout error against horizon, reward and value
calibration, and MPPI planning in simulation. Real-world deployment and dynamics finetuning
(SimDist stages 4a/4b) are out of scope for this release and have not been validated.
