model:
  arch: NanoWM-B/2
  name: NanoWM-B-2
  num_frames: 4
  n_context_frames: 1
  scheduling_mode: sequential
  num_sampling_steps: 250
  use_action: true
  action_injection:
    type: adaln
  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}