robomotic/causality-two-rooms

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

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.

Files in this repo

How to use

Load a serialized model-object checkpoint:

import torch

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

Load the full Lightning trainer checkpoint:

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.