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: placeholder — training not yet run. This card describes the model that will land here.

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.

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.

Training data

UR10e-LinearGripper-Jig-SimDist-Dataset.

Licence

MIT, following upstream SimDist.