Training the Nano World Model

conda env create -f environment.yml && conda activate nanowm

# data + output paths (or use src/configs/local/paths.yaml)
export DATASET_DIR=/path/to/dino_wm_data    # DINO-WM envs
export CSGO_DATA_DIR=/path/to/csgo          # CSGO
export RT1_DATA_ROOT=/path/to/rt1_fractal   # RT-1 (LeRobot fractal)
export RESULTS_DIR=/path/to/results         # checkpoints + logs

# i3d torchscript for periodic FID/FVD during training (one-time)
mkdir -p pretrained_models/i3d
curl -L "https://www.dropbox.com/scl/fi/c5nfs6c422nlpj880jbmh/i3d_torchscript.pt?rlkey=x5xcjsrz0818i4qxyoglp5bb8&dl=1" \
    -o pretrained_models/i3d/i3d_torchscript.pt

All runs go through src/main.py with an experiment / dataset / model triple:

# CSGO — 4x GPU, 50k steps
python src/main.py experiment=csgo dataset=game/csgo model=nanowm_l2_csgo

# RT-1 (fractal) — main run, NanoWM-B/2
python src/main.py experiment=rt1 dataset=rt1/rt1 model=nanowm_b2

# DINO-WM PushT
python src/main.py experiment=dino_wm_pusht dataset=dino_wm/pusht model=nanowm_b2

# Resume
python src/main.py experiment=csgo dataset=game/csgo model=nanowm_l2_csgo \
    resume_from_checkpoint=<path/to/ckpt>

Outputs land under ${RESULTS_DIR}/<run_dir>/: .hydra/ (composed config snapshot), checkpoints/latest/ (overwritten every 1k steps), checkpoints/across_timesteps/ (every 10k steps), and tb/ for tensorboard. Set wandb.enabled=true (+ WANDB_ENTITY / WANDB_PROJECT) for W&B logging.

PyTorch Lightning drives the loop. Validation runs every val_every_n_steps (default 1k); FID/FVD every metrics.log_every_n_train_steps (default 5k). Key knobs: experiment.training.{batch_size, max_steps, gradient_clip_norm}, experiment.diffusion.{pred_name, noise_schedule, zero_terminal_snr, snr_gamma, timestep_sampling}, experiment.infra.{mixed_precision, num_workers, compile}.

defaults:
  - model: nanowm_b2
  - dataset: dino_wm/point_maze
  - experiment: default
  - planning: base

dataset_dir:    ${oc.env:DATASET_DIR,./data}
csgo_data_dir:  ${oc.env:CSGO_DATA_DIR,./data/csgo}
vae_model_path: ${oc.env:VAE_MODEL_PATH,stabilityai/sd-vae-ft-mse}
results_dir:    ${oc.env:RESULTS_DIR,./results}

Debug composition without running: python src/main.py ... --cfg job (add --package experiment.training for one section). Common errors: ConfigCompositionException (typo in group name), MissingMandatoryValue: ckpt_path (planning needs a checkpoint), unresolvable ${oc.env:DATASET_DIR} (export the env var or pass dataset_dir= on the CLI).

python src/main.py experiment=ablation_rt1 dataset=rt1/rt1 model=nanowm_b2 \
    experiment.diffusion.pred_name=x

Reproduce any ablation arm with experiment=ablation_rt1 and the matching experiment.diffusion.* / model.action_injection.type / model= overrides. Eval numbers for these checkpoints: evaluation spec.