chomeed/dinov2-latent-wm-threading-d0

Latent world model — DINOv2 → 56-d bottleneck → dynamics (MimicGen threading_d0)

Built for RL, not for millimetres. A 2×128×28×28 feature map is 37 GB at 100k replay capacity, forcing re-encoding per sample; a 56-d latent is 11 MB and the buffer holds it directly.

Results

All at latent 56, n=40,000, 25k steps. ratio = latent prediction MSE ÷ copy-the-current-latent MSE, so 1 − ratio is the fraction of frame-to-frame latent motion the dynamics explains.

encoder	params	ratio ↓	recon ↓	expert needle / tripod / offset (mm)
conv (recommended)	0.63M	0.207	0.2009	2.21 / 1.57 / 3.02
linear	11.01M	0.249	0.2009	2.12 / 1.53 / 2.90
pooled	0.21M	0.233	0.2144	2.65 / 1.65 / 3.57

56 dimensions clears the 4 mm task requirement on expert trajectories — from a latent ~10,000× smaller than the feature maps it replaces (raw DINOv3's 602,112-value map reads 2.50 mm). The latent is never supervised on object pose: it is trained only on token reconstruction, latent dynamics and future proprioception. Pose is read out afterwards by a probe, for evaluation.

Numbers are per split because the validation set is 48% failure rollouts, where the tripod has been displaced; the pooled offsets (11.9–12.6 mm) are dominated by those and are not the regime a controller operates in.

What the encoder comparison shows

• conv beats linear on dynamics with 17× fewer parameters. Most of the frame is background; per-token weights spend capacity on it, weight-shared convolutions do not.
• pooled beats linear too (0.233 vs 0.249) with 52× fewer parameters — and pooled mean-pools the token grid, destroying spatial layout entirely. At this width, what is in the scene survives the bottleneck better than where it is.
• linear reconstructs best and predicts worst: it packs the most information into 56 dims, in a form that is harder to roll forward. For a world model that is the wrong trade.

Recommended: h1_segdepth

run	hist	dynamics ratio ↓	expert needle / tripod / offset (mm)
h1_segdepth	1	0.217	2.00 / 1.21 / 2.55
h1_tokens	1	0.214	2.14 / 1.45 / 2.83
conv56 (tokens)	2	0.207	2.21 / 1.57 / 3.02

Two findings from this pair:

History is unnecessary. Conditioning the dynamics on the current latent only costs ~nothing (0.214 vs 0.207) and slightly improves geometry. Velocity is recoverable from the action chunk and proprioception. So the RL state is a clean Markov [z(56), prop(9)] = 65 dims — no history buffer, no episode-boundary special case.

Task-relevant supervision beats token reconstruction at this width. Decoding the latent to {segmentation, depth, RGB} instead of to backbone tokens gives 2.55 mm expert offset vs 2.83 mm, at equal dynamics quality. The gain is concentrated in the TRIPOD (1.21 vs 1.45 mm), which is the weak term in offset. Pressing the latent to encode where the objects are does more than pressing it to reproduce features, ~88% of which describe background.

At 2.55 mm this 56-d latent essentially matches raw DINOv3's full 602,112-value feature map (2.50 mm) — 10,000× smaller, and never supervised on object pose.

The segdepth decoder is also a label-free reward source: its seg/depth heads run at deployment where a pose readout cannot, since the readout needs ground-truth pose to train.

Implementation

Data

MimicGen threading_d0 via LeRobot, three splits mixed into one pool (main_demos, flow_success, flow_failure), n=40,000 train / 3,072 val. Each sample is

key	shape	meaning
context	(2, 2, 256, 384)	history frames × cameras × tokens × dim
future	(2, 256, 384)	the frame at t+8, same two cameras
action	(8, 7)	the action chunk spanning t → t+8
prop / prop_fut	(9,)	proprioception at t and t+8
xyz_fut	(2, 3)	needle and tripod position at t+8, evaluation only

Both camera views (agentview, robot0_eye_in_hand) go through frozen DINOv2-small at 224×224 → 16×16 patch grid, prefix=1 (CLS only, no registers — that is a DINOv3 feature). Tokens are pre-encoded once to fp16 and cached in RAM; the backbone never runs during training, which is what makes 25k steps cheap. Tokens are z-scored per channel using statistics from the first 4,096 samples (token_mean/std in config.json).

The three losses

lc = E(context)          # (B, r, 56)   latents of the history frames
lf = E(future)           # (B, 56)      latent of the frame at t+8
lp, ph = Dy(lc, action, prop)

l_pred = mse(lp, lf.detach())            # dynamics — stop-grad on the target
l_rec  = <decoder loss on lf>            # keeps information inside the bottleneck
l_prop = mse(ph, prop_fut_normalised)    # future proprioception, an auxiliary anchor
loss   = l_pred + 1.0 * l_rec + l_prop

lf.detach() is the important detail. Without the stop-grad, the encoder can minimise l_pred by making all latents identical — the dynamics term alone has a trivial collapse solution. Detaching means the prediction loss can only be reduced by improving the dynamics, never by degrading the representation; l_rec then independently pushes back on how much the latent may throw away.

AdamW lr=1e-3, wd=1e-4, OneCycle over 25,000 steps (pct_start=0.05), batch 128, grad-norm clipped to 1.0.

l_rec: the two variants this repo compares

• tokens (Dec) — 56 → 512 → v·64·4·4, two ConvTranspose2d stages back to the 16×16 grid, MSE against the normalised DINOv2 tokens. Generic: reproduce what the backbone saw.

• segdepth (SegDec) — 56 → 512 → v·64·7·7, then Conv2d + PixelShuffle(2) blocks doubling 7→14→28→56, and three 1×1 heads for segmentation (4 classes), depth, and RGB:

  cross_entropy(seg, seg_gt, weight=CW) + 5.0 * l1(depth, depth_gt) + l1(rgb, rgb_gt)
  

  Targets are MuJoCo sim-state replays rendered through EGL, cached as uint8/fp16. Segmentation is downsampled by majority vote with rare-class priority (for cl in (3,0,1,2)), never by averaging — averaging class ids invents labels that do not exist. Depth and RGB use average pooling, which is correct for continuous targets.

  This is still label-free with respect to the task: seg/depth/RGB come from the simulator's renderer, not from object-pose annotations. That matters because it is what a pose readout cannot claim — the readout needs ground-truth pose to train, so it cannot supply a deployment-time reward.

Encoders

All three end in LayerNorm(56), which is load-bearing (see below).

• conv — Conv2d(384→128, s2) → GELU → Conv2d(128→64, s2) → GELU → Linear, applied per camera over the 16×16 grid, then both cameras concatenated. 0.63M params.
• linear — one Linear(2·256·384 → 56). Every token gets its own weights. 11.01M params.
• pooled — mean over tokens, then an MLP. Spatial layout is destroyed; this is the control for whether layout matters at 56 dims. 0.21M params.

Dynamics

z_next = z_t + MLP([z_ctx.flatten(), action_chunk(8×7), proprio(9)])

A residual predictor: the network outputs a delta, not an absolute latent, so the identity map is free and it only has to learn what changes. r=hist slices zc[:, -r:], so hist=1 conditions on the current latent only (input 56+56+9 = 121) and hist=2 also sees the previous one (112+56+9 = 177). A second small head predicts future proprioception from the conditioning alone.

Why LayerNorm on the latent

ratio = pred_mse / copy_mse, where copy_mse is the error of simply reusing the current latent. Both terms scale with the latent's magnitude, so an unconstrained encoder can inflate the scale and drive the ratio down while predicting nothing better. Observed directly during development: the linear encoder grew copy from 4,871 → 74,732 in 5,000 steps while reconstruction stayed flat. The LayerNorm pins the scale and makes the ratio comparable across runs.

Evaluation

The pose probe is trained after the world model, on frozen latents: [z(56), prop(9)] → 512 → 256 → 6, smooth-L1 on z-scored (needle_xyz, tripod_xyz), 6,000 steps. Predictions are de-normalised with target_mean/std and reported in millimetres. offset = needle − tripod is the decision metric, since relative pose is what a controller consumes.

Results are reported per split. Validation is 48% failure rollouts where the tripod has been knocked out of place, so pooled numbers (11.9–12.6 mm) describe a regime the controller never operates in; expert is the number to read.

Contents

{conv,linear,pooled}56/ (hist=2, token target) and {h1_tokens,h1_segdepth}/ (hist=1) — each holds encoder, dynamics, decoder, pose_probe as safetensors plus a config.json carrying every normalisation statistic (prop_mean/std, target_mean/std, token_mean/std) and the measured pred_mse / copy_mse / ratio. measurements/ holds the training logs.

Every statistic needed to run these checkpoints is in config.json — normalise tokens with token_* before the encoder, proprioception with prop_* before the dynamics, and de-normalise probe output with target_*. (Shipping a checkpoint without its normalisation affine makes it unusable; that mistake cost real time on an earlier release and is not repeated here.)

Minimal forward pass

tok = (dinov2(images).last_hidden_state[:, 1:] - token_mean) / token_std   # (B,2,256,384)
z   = encoder(tok)                                                        # (B,56)
state = torch.cat([z, (prop - prop_mean) / prop_std], -1)                  # (B,65)  <- RL state
z_next, prop_next = dynamics(z[:,None], action_chunk, (prop-prop_mean)/prop_std)
xyz = pose_probe(z, (prop-prop_mean)/prop_std) * target_std + target_mean  # mm, eval only

Caveats

• Single seed per encoder; the conv-vs-linear gap (0.207 vs 0.249) is larger than the drift between checkpoints, but has not been seed-repeated.
• Only an 8-step horizon was trained and measured. MPC/SVG would use shorter horizons, where error is lower.
• Action-gradient quality through this dynamics model has not been measured yet — worth checking before relying on it for policy gradients.