File size: 2,017 Bytes
088b1b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | ---
library_name: pytorch
tags:
- video-generation
- world-model
- diffusion
- diffusion-forcing
- flow-matching
- ablation
datasets:
- lerobot/fractal20220817_data
---
# NanoWM-B/2 · RT-1 · Ablation: pred_name = flow
Checkpoint from the prediction-target ablation on RT-1 fractal
(flow-matching arm).
## Run identity
- **wandb**: https://wandb.ai/better_guidance/nano-world-model-ablation/runs/uw3i6shn
- **launcher**: `src/scripts/train/ablation/pred_flow.sh`
- **collection**: https://huggingface.co/collections/knightnemo/nano-world-model
## Training setup
| Key | Value |
|---|---|
| Architecture | NanoWM-B/2 (12 layers, d=768, patch=2, 158.6M params) |
| Dataset | RT-1 fractal (`lerobot/fractal20220817_data`) |
| Frames × resolution | 4 × 256² → 4 × 32² latents (SD-VAE) |
| Context frames | 1 (sequential / self-forcing scheduling) |
| Action injection | additive (7-dim continuous) |
| Steps | 50,000 |
| Batch | 8/GPU × 8 × H20 = 64 effective |
| Optimizer | AdamW, lr 1e-4, wd 0.01, warmup 1000, grad clip 0.1 after 20k |
| Precision | bf16-mixed (params fp32), VAE fp32, `torch.compile` on |
| Seed | 3407 |
## Diffusion setup
| Key | Value |
|---|---|
| pred_name | **flow** |
| noise_schedule | `squaredcos_cap_v2` |
| zero_terminal_snr | false |
| timestep_sampling | logit_normal (SD3-style, μ=0, σ=1) |
| snr_gamma | 0.0 |
| diffusion_steps | 1000 train · 250 sample |
| history_stabilization_level (inference) | 0.02 |
## Loading
```bash
git clone git@github.com:simchowitzlabpublic/nano-world-model.git
cd nano-world-model
huggingface-cli download knightnemo/nanowm-b2-rt1-abl-pred-flow-50k --local-dir ./ckpt
```
```python
import sys
from omegaconf import OmegaConf
from safetensors.torch import load_file
sys.path.insert(0, "src")
from models import get_models
cfg = OmegaConf.load("ckpt/config.yaml")
cfg.experiment.infra.compile = False
model = get_models(cfg).eval()
state_dict = load_file("ckpt/model.safetensors")
model.load_state_dict(state_dict, strict=True)
```
|