Initial upload: scale_s2

README.md

@@ -0,0 +1,71 @@
+---
+library_name: pytorch
+tags:
+- video-generation
+- world-model
+- diffusion
+- diffusion-forcing
+- ablation
+- scaling
+datasets:
+- lerobot/fractal20220817_data
+---
+
+# NanoWM-S/2 · RT-1 · Ablation: model scale = S/2
+
+Scale ablation arm on RT-1 fractal. All other axes (pred_target=v, additive
+action injection, cosine + ZTSNR schedule, 50k steps) match the B/2 reference,
+so this row directly isolates the effect of backbone capacity.
+
+## Run identity
+- **wandb**: https://wandb.ai/better_guidance/nano-world-model-ablation/runs/qdwuh9oe
+- **launcher**: `src/scripts/ablation/scale_s2.sh`
+- **collection**: https://huggingface.co/collections/knightnemo/nano-world-model
+
+## Training setup
+| Key | Value |
+|---|---|
+| Architecture | **NanoWM-S/2** (12 layers, d=384, patch=2, ~39.8M 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 | v |
+| noise_schedule | `squaredcos_cap_v2` (cosine) |
+| zero_terminal_snr | true |
+| timestep_sampling | logit_normal (SD3-style, μ=0, σ=1) |
+| snr_gamma | 5.0 (Min-SNR loss weighting) |
+| diffusion_steps | 1000 train · 250 DDIM sample |
+| history_stabilization_level (inference) | 0.02 |
+
+## Loading
+
+```bash
+git clone git@github.com:knightnemo/nano-world-model.git
+cd nano-world-model
+huggingface-cli download knightnemo/nanowm-s2-rt1-abl-scale-s2-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)
+```

config.yaml

@@ -0,0 +1,109 @@
+model:
+  arch: NanoWM-S/2
+  name: NanoWM-S-2
+  num_frames: 4
+  n_context_frames: 1
+  scheduling_mode: sequential
+  num_sampling_steps: 250
+  use_action: true
+  action_injection:
+    type: additive
+  causal: true
+  image_size: 256
+  latent_size: 32
+  extras: 1
+  num_classes: 1000
+dataset:
+  loader:
+    n_rollout: null
+    data_path_train: null
+    data_path_val: null
+    split_ratio: 0.9
+    validation_size: 32
+    normalize_state: false
+    normalize_action: true
+    train_slice_mode: random
+    val_slice_mode: exhaustive
+    stride: 1
+    random_seed: 42
+    validation_fixed_subset_path: null
+    validation_fixed_subset_size: null
+    validation_fixed_subset_seed: 42
+    data_path: lerobot/fractal20220817_data
+    root: ${oc.env:RT1_DATA_ROOT,/wuji-vepfs/wuji-il/huangsiqiao/nano-world-model-data/rt1_fractal}
+    image_key: observation.images.image
+  name: rt1
+  frame_interval: 1
+  spec:
+    action_dim: 7
+experiment:
+  name: train
+  tasks:
+  - training
+  resume_from_checkpoint: null
+  pretrained: null
+  training:
+    optimizer:
+      lr: 0.0001
+      weight_decay: 0.01
+      lr_warmup_steps: 1000
+    max_steps: 50000
+    batch_size: 8
+    gradient_accumulation: 1
+    gradient_clip_norm: 0.1
+    gradient_clip_start_step: 20000
+    log_every: 100
+    val_every_n_steps: 1000
+    checkpointing:
+      across_timesteps:
+        every_n_train_steps: 10000
+        save_top_k: -1
+        save_on_train_epoch_end: true
+        save_weights_only: false
+        filename: '{epoch}-{step}'
+      latest:
+        every_n_train_steps: 1000
+        save_top_k: 1
+        save_on_train_epoch_end: false
+        save_weights_only: false
+        filename: latest-{epoch}-{step}
+  evaluation:
+    validation_size: 32
+    save_videos: true
+    metrics:
+      evaluate: true
+      log_every_n_train_steps: 5000
+      buffer_size: 32
+      max_batchsize: 2
+      i3d_model_path: ${oc.env:PRETRAINED_MODELS_DIR,pretrained_models}/i3d/i3d_torchscript.pt
+  diffusion:
+    noise_schedule: squaredcos_cap_v2
+    diffusion_steps: 1000
+    pred_name: v
+    mode: diffusion_forcing
+    snr_gamma: 5.0
+    zero_terminal_snr: true
+    timestep_sampling: logit_normal
+    logit_normal_mean: 0.0
+    logit_normal_std: 1.0
+    history_stabilization_level: 0.02
+  infra:
+    mixed_precision: true
+    vae_precision: fp32
+    gradient_checkpointing: false
+    num_workers: 16
+    compile: true
+    seed: 3407
+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}
+logger:
+  name: wandb
+  save_dir: ${hydra:runtime.output_dir}/tb
+  logger_name: nanowm
+wandb:
+  enabled: true
+  entity: ${oc.env:WANDB_ENTITY,null}
+  project: nano-world-model-ablation
+  mode: ${oc.env:WANDB_MODE,online}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8169cc27064d53c4f210271a0273d243898ff2d91a71a0220c7ba9bf60af270
+size 159140872