| # 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: 60 # latent frame count; dataset loads (21-1)*4+1=81 video frames | |
| load_raw_video: true | |
| sample_strategy: random # 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) | |
| - 60 # 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 | |
| latent_channels: 16 | |
| # Training hyperparams | |
| num_train_timestep: 1000 | |
| timestep_shift: 5.0 | |
| guidance_scale: 1.0 | |
| teacher_forcing: true | |
| noise_augmentation_max_timestep: 0 | |
| generator_ckpt: exps/dense_short/checkpoint_model_022000/model.pt | |
| lr: 2.0e-5 | |
| 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: 500 | |
| # ema_weight: 0.999 | |
| # ema_start_step: 500 | |
| gc_interval: 100 | |
| wandb_host: "https://api.wandb.ai" | |
| wandb_key: | |
| wandb_entity: "" | |
| wandb_project: "" | |
| negative_prompt: "" | |
| # FSDP wrap strategy | |
| generator_fsdp_wrap_strategy: size | |
| text_encoder_fsdp_wrap_strategy: size |