Trained model checkpoints from the project on improving counter-factual
reasoning in LeWorldModel (LeWM) via slot-attention / DINOSAUR-style
binding / C-JEPA-style masking.
Companion docs (in the GitHub code repo):
SOLUTIONS_TRIED.md — full method index with sources
EP100_N50_GROUND_TRUTH.md — eval history per phase
Cube (OGBCube-v0, main task, n=20)
File
Method
Source idea
baseline
near_ood
CF6_no_far
cube/longtrain_baseline_ep20.ckpt
LongTrain
vanilla LeWM
65
10
55.8
cube/visualrand_s43_ep20.ckpt
VisualRand
RandConv (Lee 2020)
65
10
55.8
cube/orid_inv001_vr_s42_ep20.ckpt
ORID + Inverse Dyn
ICM (Pathak 2017)
65
10
58.3
cube/orid_foreground_s42_ep20.ckpt
ORID + Foreground mask
self / motion mask
65
15
56.7
cube/orid_posexy_low_s42_ep8.ckpt
ORID + Pose Grounding (priv label)
self
65
15
56.7
cube/fair_slot_v3_s42_ep3_INTERIM.ckpt
Fair-Slot-LeWM v3 (DINOSAUR + Hungarian)
DINOSAUR (Seitzer 2023) + SlotFormer
INTERIM
INTERIM
INTERIM
TwoRoom (2D toy, n=20)
File
Method
Source idea
baseline
CF6
tworoom/longtrain_baseline_ep100.ckpt
LongTrain
vanilla LeWM
0.86-0.90
~0.67
tworoom/visualrand_ep53.ckpt
VisualRand
RandConv
1.00
0.79 (best fair)
tworoom/cfcons_ep33.ckpt
CF-Consistency
self-consistency JEPA (Garrido 2023)
0.85
0.70
tworoom/factorbottleneck_ep33.ckpt
Latent Factor Bottleneck
beta-VAE
0.85
0.70
tworoom/actiongeo_ep31.ckpt
Action-Geometry Aux
GeoAux-style
0.85
0.70
tworoom/slot_tworoom_s44_ep24_FAILED.ckpt
SlotJEPA (FAILED)
SlotFormer (Wu 2023) + DynSlotAttn
0.45
0.40
Architecture
All checkpoints share the LeWM ViT-Tiny encoder (~5.5M params, 192d, 12
layers). Total model size:
import torch
ckpt = torch.load('cube/longtrain_baseline_ep20.ckpt', map_location='cpu')
# ckpt is a JEPA model object; eval pipeline lives in
# delta_lewm/counterfactual_eval/eval_counterfactual_cube.py
Citation
@article{nam2026cjepa,
title = {Causal-JEPA: Learning World Models through Object-Level Latent Interventions},
author = {Nam, Heejeong and Le Lidec, Quentin and Maes, Lucas and LeCun, Yann and Balestriero, Randall},
journal = {arXiv preprint arXiv:2602.11389},
year = {2026},
}
Phase 8 (2026-04-28): Fair-Slot v3 + Lite-CJEPA (S/M/L) — final
Lite-CJEPA-S (Locatello SlotAttn + masked-slot pred + recon)
cube/lite_cjepa_M_ep20.ckpt
9.7M
d=6 h=8 m=1024 (5.5M)
Lite-CJEPA-M (medium predictor)
cube/lite_cjepa_L_ep20.ckpt
17.7M
d=12 h=16 m=2048 (10.4M)
Lite-CJEPA-L (predictor matched to LeWM total) — uploaded after training
Cube CF6 results (n=20, seed=42)
Run
base
clr_ood
swap
floor
sm
lg
near_init
CF6_no_far
LongTrain ep20 (baseline)
65
65
65
65
65
65
10
55.83
v3 ep5
70
65
65
65
70
65
10
56.67
v3 ep15
60
65
65
65
60
70
15
56.67
v3 ep20
65
60
65
60
65
65
10
54.17
lcS ep17
65
65
65
65
65
65
10
55.83
lcS ep20
65
65
65
65
65
65
10
55.83
lcM ep17
65
65
65
65
65
65
10
55.83
lcL ep17
65
65
65
65
65
65
10
55.83
Conclusion
near_ood_init = 10% is a hard ceiling across all 4 architectures, predictor sizes (1.1M ~ 10.4M), and training epochs (5-20).
Slot architecture / DINOSAUR feature reconstruction / cjepa masked-slot prediction do not break this geometric OOD wall on cube_single_expert.
Hypothesis: ViT absolute positional embedding hard-codes object location into slot tokens; when cubes appear at OOD start positions the predictor cannot extrapolate in slot-coordinate space.
1import torch
2from huggingface_hub import hf_hub_download
34ckpt = hf_hub_download(5 repo_id="dzyy123/lewm-counterfactual-checkpoints",6 filename="cube/lite_cjepa_S_ep20.ckpt",7)8model = torch.load(ckpt, map_location="cpu", weights_only=False).eval()9print(type(model).__name__,sum(p.numel()for p in model.parameters()))
For lite_cjepa_*, you must add /path/to/lite-cjepa/src to sys.path so that
the model, encoder, decoder, predictor modules can be unpickled
(see lite-cjepa GitHub repo for source).
