---
license: mit
library_name: pytorch
pipeline_tag: reinforcement-learning
tags:
- robotics
- reinforcement-learning
- world-model
- causal-representation-learning
- stable-worldmodel
---

# 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:

```bash
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:

```bash
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.