You can deploy this model and start earning money today!

robomotic/
causality-two-rooms

Name: causality-two-rooms
Availability: InStock

pytorch

robotics

reinforcement-learning

world-model

causal-representation-learning

stable-worldmodel

mit

Views

No views yet

robomotic/causality-two-rooms

LeWM checkpoints trained on the confounded Glitched Hue TwoRoom dataset for causal world-model experiments. The goal is to test whether the model learns the true teleport mechanism or the spurious background-hue correlation.

Model description

Architecture: LeWM / JEPA-style world model with an autoregressive predictor
Domain: swm/GlitchedHueTwoRoom-v1
Framework: PyTorch + Lightning
Repository: https://github.com/epokhcs/stable-worldmodel
Upload generated (UTC): 2026-04-10 08:16:57

Training data

Dataset path: /home/robomotic/.stable_worldmodel/glitched_hue_tworoom.h5
Dataset size: 9.09 GB
Episodes: 10,000
Frames: 609,539
Pixel tensor shape: (609539, 224, 224, 3)
Teleport events: 1,593

The dataset was collected with:

python scripts/data/collect_glitched_hue.py num_traj=10000 seed=3072 world.num_envs=10

Training procedure

The checkpoints in this repo come from the 5-epoch LeWM training run used in the causality experiment. The run completed successfully after resuming from the last full trainer checkpoint.

Command family:

python scripts/train/lewm.py     data=glitched_hue_tworoom     trainer.max_epochs=5     num_workers=1     loader.num_workers=1     loader.persistent_workers=False

Key hyperparameters

Parameter	Value
`trainer.max_epochs`	`5`
`trainer.accelerator`	`gpu`
`trainer.precision`	`bf16`
`loader.batch_size`	`128`
`loader.num_workers`	`1` for the resumed run
`optimizer.lr`	`5e-5`
`wm.history_size`	`3`
`wm.num_preds`	`1`
`wm.embed_dim`	`192`
`loss.sigreg.weight`	`0.09`
`data.dataset.frameskip`	`5`
`config.yaml`	included in the repo root

Epoch metrics (logged to W&B / Lightning)

The table below summarizes the epoch-end losses extracted from the local training logs. The raw values are also included as metrics/epoch_metrics.json and metrics/epoch_metrics.csv.

Epoch	Global step	fit/loss	fit/pred_loss	fit/sigreg_loss	validate/loss	validate/pred_loss	validate/sigreg_loss
1	3046	2.676974	0.005099	29.750000	2.766938	0.009671	30.635967
2	6092	1.143889	0.436858	7.843750	1.173404	0.444629	8.096438
3	9138	0.673702	0.310420	4.031250	0.746696	0.313994	4.807232
4	12184	0.526915	0.163634	4.031250	0.601880	0.166302	4.839801
5	15230	0.496008	0.156165	3.765625	0.520261	0.127601	4.362802

Files in this repo

File	Purpose	Size
`checkpoints/lewm_epoch_1_object.ckpt`	Serialized model object checkpoint	69.00 MB
`checkpoints/lewm_epoch_2_object.ckpt`	Serialized model object checkpoint	69.00 MB
`checkpoints/lewm_epoch_3_object.ckpt`	Serialized model object checkpoint	69.00 MB
`checkpoints/lewm_epoch_4_object.ckpt`	Serialized model object checkpoint	69.00 MB
`checkpoints/lewm_epoch_5_object.ckpt`	Serialized model object checkpoint	69.00 MB
`checkpoints/lewm_weights.ckpt`	Full Lightning trainer checkpoint	206.75 MB
`config.yaml`	Hydra config used for the run	1.20 KB
`metrics/epoch_metrics.json`	Raw epoch metrics extracted from local logs	small
`metrics/epoch_metrics.csv`	Tabular epoch metrics for spreadsheets / plotting	small

How to use

Load a serialized model-object checkpoint:

python

import torch

model = torch.load('checkpoints/lewm_epoch_5_object.ckpt', map_location='cpu')
model.eval()

Load the full Lightning trainer checkpoint:

python

import torch

checkpoint = torch.load('checkpoints/lewm_weights.ckpt', map_location='cpu')
print(checkpoint.keys())

Intended uses

Reproducing the causal disentanglement experiment in research/runme.md
Running the Step 3 causal AAP analysis with research/glitched_hue_experiment.py
Comparing epoch-wise world-model checkpoints during training

Limitations

These checkpoints are research artifacts, not production control policies.
Performance is specific to the Glitched Hue TwoRoom environment and the confounded blue/green data collection procedure.
The object checkpoints are convenient for inspection, while the full trainer checkpoint is the correct file for resuming optimization.