RubetekRobotics/UR10e-LinearGripper-Jig-SimDist-WorldModel

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

Latent world model pretrained in simulation with Simulation Distillation (SimDist) (arXiv:2603.15759, RSS 2026; code 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.