world-model

Summary In this post, I train LeWM without SIGReg and compare its embeddings with normal runs. Q: What’s the purpose of this experiment? I want to see what SIGReg does. To do this, I compare the embeddings of models trained with and without SIGReg, and look for representation collapse in the model without it. The non-SIGReg models is trained with the same setup as my baseline (image size 112, 3 epochs), except that the SIGReg weight ($\lambda) is set to 0. ...

LeWM Push-T: My Second Experiment vs Paper (arXiv 2603.19312v3) This run fixes the two main differences from the first experiment: it trains for 10 epochs instead of 3, and it uses 224×224 images instead of 112×112, for both training and eval. Legend: ✓ = matches · ≈ = close / likely same · ⚠️ = differs, likely affects results · — = not stated in paper 1. Training parameters Mine Paper Epochs 10 10 ✓ Image size 224×224 224×224 ✓ Encoder ViT-tiny, patch 14, 12 layers, 3 heads, hidden 192, from scratch ViT-tiny, patch 14, 12 layers, 3 heads, hidden 192, from scratch ✓ Embedding [CLS] → MLP projector (hidden 2048) + BatchNorm, dim 192 [CLS] → 1-layer MLP + BatchNorm, dim 192 ≈ Predictor transformer, 6 layers, 16 heads, dropout 0.1, AdaLN action conditioning, ~10M params transformer (“ViT-S”), 6 layers, 16 heads, 10% dropout, AdaLN, ~10M params ✓ History size 3 3 ✓ Sub-trajectory 4 frames + 4 blocks of 5 actions 4 frames + 4 blocks of 5 actions ✓ Frameskip 5 5 ✓ Batch size 128 128 ✓ SIGReg λ 0.09 0.1 (λ ablation peaks near 0.09) ≈ SIGReg projections 1024 1024 ✓ SIGReg knots 17 T nodes in [0.2, 4] (count not stated; insensitive) — Optimizer / LR / WD AdamW / 5e-5 / 1e-3 — — Precision / grad clip bf16 / 1.0 — — Train / val split 0.9 / 0.1 — — Dataset librakevin/lewm-pusht DINO-WM Push-T, 20,000 expert episodes ≈ (not verified) Training seeds 1 (seed 3072) 3 ⚠️ The run took 6 h 40 min on one RTX 5090: 13,933 steps per epoch and 139,330 steps in total. The full training log is in train.log. ...

LeWM Push-T: My Experiment vs Paper (arXiv 2603.19312v3) Legend: ✓ = matches · ≈ = close / likely same · ⚠️ = differs, likely affects results · — = not stated in paper 1. Training parameters Mine Paper Epochs 3 10 ⚠️ Image size 112×112 224×224 ⚠️ Encoder ViT-tiny, patch 14, 12 layers, 3 heads, hidden 192, from scratch ViT-tiny, patch 14, 12 layers, 3 heads, hidden 192, from scratch ✓ Embedding [CLS] → MLP projector (hidden 2048) + BatchNorm, dim 192 [CLS] → 1-layer MLP + BatchNorm, dim 192 ≈ Predictor transformer, 6 layers, 16 heads, dropout 0.1, AdaLN action conditioning, ~10M params transformer (“ViT-S”), 6 layers, 16 heads, 10% dropout, AdaLN, ~10M params ✓ History size 3 3 ✓ Sub-trajectory 4 frames + 4 blocks of 5 actions 4 frames + 4 blocks of 5 actions ✓ Frameskip 5 5 ✓ Batch size 128 128 ✓ SIGReg λ 0.09 0.1 (λ ablation peaks near 0.09) ≈ SIGReg projections 1024 1024 ✓ SIGReg knots 17 T nodes in [0.2, 4] (count not stated; insensitive) — Optimizer / LR / WD AdamW / 5e-5 / 1e-3 — — Precision / grad clip bf16 / 1.0 — — Dataset librakevin/lewm-pusht DINO-WM Push-T, 20,000 expert episodes ≈ (not verified) Training seeds 1 (seed 3072) 3 ⚠️ 2. Eval parameters Mine Paper Episodes 50 50 (same 50 for all models) ✓ Eval budget 50 steps 50 steps ✓ Goal offset 25 steps 25 steps ✓ CEM samples 300 300 ✓ CEM iterations 30 30 (Push-T) ✓ CEM top-k 30 30 ✓ CEM initial variance 1.0 1 ✓ Planning horizon 5 (= 25 env steps) 5 (= 25 env steps) ✓ Receding horizon 5 (execute full plan) 5 (execute full plan) ✓ Eval image size 112 224 ⚠️ Eval seed 42 (default in config/eval/pusht.yaml) — — 3. Eval Result (Push-T success rate) All 10 seeds together (same epoch-3 checkpoint, 500 episodes in total): ...

How LeWorldModel prevents representation collapse with SIGReg, with interactive visualizations of one training step, the characteristic-function test, and CEM planning.