dzyy123/lewm-counterfactual-checkpoints

LeWM Counter-Factual Reasoning Checkpoints

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:

• LeWM family (longtrain / visualrand / cfcons / factor / actiongeo / orid / fair-slot-v3): ~17M params
• SlotJEPA family: ~7M params (smaller predictor)

Loading

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

ckpt	params	predictor	description
cube/v3_ep20.ckpt	17.3M	LeWM AR + Hungarian	Fair-Slot-LeWM v3 (DynamicSlot K=4, Hungarian + DINOSAUR recon)
cube/lite_cjepa_S_ep20.ckpt	7.5M	d=3 h=4 m=512 (1.1M)	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.

See SOLUTIONS_TRIED.md §10 for full ablation table and discussion.

Quick load

import torch
from huggingface_hub import hf_hub_download

ckpt = hf_hub_download(
    repo_id="dzyy123/lewm-counterfactual-checkpoints",
    filename="cube/lite_cjepa_S_ep20.ckpt",
)
model = torch.load(ckpt, map_location="cpu", weights_only=False).eval()
print(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).