---
license: bsd-3-clause
tags:
  - world-model
  - dynamics-model
  - robotics
  - manipulation
datasets:
  - villekuosmanen/bin_pick_pack_coffee_capsules
---

# Visual World Model — bin-pick-pack (proprio + visual latents)

MLP-based world model (`SystemDynamicsEnsemble`) trained on the bin-pick-pack-coffee-capsules manipulation dataset. Predicts next proprioceptive state (17D) and next visual latent (32D) given a history of states, actions, and visual latents.

## Architecture

- **Type**: MLP ensemble (2 heads)
- **State dim**: 17 (7 joint positions + 10 EEF pose as xyz/rot6d/gripper)
- **Action dim**: 17 (same decomposition)
- **Visual dim**: 32 (LAM-encoded visual latents from front camera)
- **History horizon**: 2
- **Forecast horizon**: 1
- **Checkpoint size**: 1.9 MB

## Training

- **Dataset**: [villekuosmanen/bin_pick_pack_coffee_capsules](https://huggingface.co/datasets/villekuosmanen/bin_pick_pack_coffee_capsules) — 47865 frames, 200 episodes
- **Visual latents**: Precomputed from fine-tuned LAM encoder ([pravsels/lam-binpack-finetune](https://huggingface.co/pravsels/lam-binpack-finetune))
- **Split**: 35387 train / 12078 val sequences (val_ratio=0.25, seed=0)
- **Epochs**: 50
- **Batch size**: 64
- **Learning rate**: 3e-4 (cosine schedule, min_lr=3e-5)
- **Final train loss**: 0.09442
- **Final val loss**: 0.09916
- **Visual loss**: 0.07207
- **W&B**: [pravsels/binpack-world-model/runs/2pq0n2mx](https://wandb.ai/pravsels/binpack-world-model/runs/2pq0n2mx)

## Files

| File | Description |
|------|-------------|
| `best.pt` | Best checkpoint (epoch 50) |
| `config.yaml` | Training configuration (Isambard) |

## Checkpoint format

```python
checkpoint = torch.load("best.pt", map_location="cpu")
# checkpoint["model_state_dict"] -> SystemDynamicsEnsemble.load_state_dict()
# checkpoint["epoch"], checkpoint["train_loss"], checkpoint["val_loss"], etc.
```

## Integrity

```
sha256: 050dfaffa2c98ff112d6a0d2eba738328bac8b3934863bfecff59d62bd2d2410  best.pt
```

Verified by running `sha256sum` twice on the source file.

## Usage

```python
from rsl_rl.offline.offline_world_model_trainer import build_system_dynamics_model
import torch

model = build_system_dynamics_model(
    state_dim=17, action_dim=17, visual_dim=32,
    ensemble_size=2, history_horizon=2, device="cpu",
)
ckpt = torch.load("best.pt", map_location="cpu")
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

# Single-step prediction
# state_hist: (1, 2, 17), action_hist: (1, 2, 17), visual_hist: (1, 2, 32)
state_pred, _, _, _, _, _, visual_pred = model(state_hist, action_hist, visual_hist)
```