push_cube_ar.yaml

# ACWM-DiT-AR: autoregressive shortcut-forcing world model on push_cube.
# Latent (Wan-VAE) + causal DiT-S + shortcut forcing (velocity-prediction). Few-step AR rollout.
model_name: "ShortcutDiT"
dynamics_class: "ShortcutForcing_WM"

model_config:
  # --- latent DiT-S (causal) ---
  in_channels: 16
  patch_size: 2
  action_compress_rate: 4
  max_frames: 37
  dim: 768
  num_layers: 10
  num_heads: 12
  action_dropout_prob: 0.0
  temporal_causal: true        # AR => causal temporal attention (wrapper forces it too)
  use_flash_attn: true
  action_conditioning: "adaLN"
  # --- Wan-VAE ---
  vae_name: "WanVAE"
  vae_config:
    - "Wan2.1_VAE.pth"         # WAN_VAE_PATH env overrides
  temporal_compress_rate: 4
  # --- shortcut forcing (Dreamer4) ---
  k_max: 8                     # noise/step grid; step sizes d in {1,1/2,1/4,1/8}
  self_fraction: 0.25          # fraction of batch using the self-consistency (bootstrap) loss
  bootstrap_start: 5000        # start the consistency term after this many steps
  eval_d: 0.25                 # rollout step size at eval => 4 steps/frame

dataset:
  name: "push_cube"
  seq_len: 37
  obs_shape: [3, 240, 240]
  train_size: 40000
  test_cuts: 5

training:
  batch_size: 8                # H100 smoke (w/ bootstrap): bs8=38%, bs16=70%. global 8x4=32
  learning_rate: 1e-4
  num_epochs: 1000
  total_steps: 100000
  num_workers: 8
  grad_clip: 1.0
  log_freq: 10
  val_freq: 1000
  checkpoint_freq: 2000

wandb:
  project: "AC-WM-Phys"
  run_name: "AR_ShortcutDiT_S_push_cube"
  api_key: "62da90010e5c8cc94a66361396c57cea8c2c1e21"

distributed:
  use_fsdp: false