# WorldMem diffusion training config
# Action + PRoPE conditioning, no text conditioning, CFG on actions only

trainer: diffusion
causal: true
mixed_precision: true
seed: 0

# Data
# CSV columns: video_path, action_path  (WorldMem format)
data_path: data/train_wm_gamecraft/v2/worldmem_train_v2_shuffle.csv
num_frames: 28          # latent frame count; dataset loads (21-1)*4+1=81 video frames
load_raw_video: true
sample_strategy: first  # overfit: always start from SKIP_FRAMES; use "random" for normal training

# CFG: probability of dropping action conditioning (zeros actions + identity viewmats)
proportion_empty_prompts: 0.1

# Model
# 81 video frames @ 640x352 -> 21 latent frames @ 80x44 (WAN 4x temporal, 8x spatial)
image_or_video_shape:
- 1        # batch_size placeholder (overridden at runtime)
- 28       # latent frames
- 16       # latent channels
- 44       # latent height  (352/8)
- 80       # latent width   (640/8)

model_kwargs:
  timestep_shift: 5.0
  local_attn_size: -1
  sink_size: 0
  use_action_encoder: true
  action_dim: 25
  action_hidden_dim: 128
  latent_channels: 16

# Training hyperparams
num_train_timestep: 1000
timestep_shift: 5.0
guidance_scale: 1.0
teacher_forcing: true
# noise_augmentation_max_timestep: 100
noise_augmentation_max_timestep: 0

lr: 1.0e-4
beta1: 0.9
beta2: 0.999
weight_decay: 0.01
batch_size: 1
gradient_checkpointing: true
num_frame_per_block: 4
sharding_strategy: hybrid_full

# Logging / saving
no_save: false
no_visualize: true
log_iters: 300
ema_weight: 0.999
ema_start_step: 200
gc_interval: 100
wandb_host: "https://api.wandb.ai"
wandb_key: WANDB_KEY
wandb_entity: "zhenhao-yang11-the-university-of-hong-kong"
wandb_project: worldmem_diffusion
negative_prompt: ""

# FSDP wrap strategy
generator_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size