knightnemo/nanowm-b2-rt1-abl-pred-flow-50k

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

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

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)