README.md

---
license: mit
library_name: jax
tags:
- robotics
- world-models
- jax
- flax
- isaac-sim
- manipulation
- simdist
---

# UR10e Linear Gripper — Jig / Bottom Enclosure — SimDist World Model

Latent world model pretrained in simulation with **Simulation Distillation (SimDist)**
([arXiv:2603.15759](https://arxiv.org/abs/2603.15759), RSS 2026; code
[CLeARoboticsLab/simdist](https://github.com/CLeARoboticsLab/simdist), MIT).

**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).

## Training

```
python scripts/train_model.py \
  model=manipulator_world_model system=omnireset_ur10e \
  data.dataset_name=simdist_merged data.num_train_workers=32 data.num_test_workers=4 \
  training.batch_size=256 training.max_steps=90000 training.warmup_steps=10000 \
  training.decay_steps=90000 training.eval_interval=2500 \
  training.grad_clip_norm=1.0 model.encoder.latent_norm=true \
  checkpoint.enabled=True checkpoint.max_to_keep=5 \
  run_name=simdist_ur10e_24m_v2
```

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.

## Licence

MIT, following upstream SimDist.