actions:
  target_fps: 20
  valid_keys:
  - W
  - A
  - S
  - D
  - Q
  - E
  - Space
  - LShiftKey
  - LControlKey
wandb:
  entity: null
  project: mira-world-model
  name: world_model-260711-2008
  group: null
  mode: disabled
run:
  seed: 28
  steps: 10001
  batch_size: 2
  checkpoint_every: 500
  checkpoint_keep_recent: 3
  checkpoint_keep_permanent_every: 2500
  log_every: 50
  output_dir: /home/hugohernandez/runs/wm1b_psd10k
  compile: false
  continue_from: /home/hugohernandez/runs/wm1b_psd10k/checkpoint-1500
  finetune_from: null
dataloader:
  num_workers: 6
  shuffle_buffer_size: 100
validation:
  batch_size: null
  val_first: true
  val_n_samples: 256
  val_every: 1000
  downstream_val_every: 2500
optim:
  optimizer:
    lr: 3.0e-05
    betas:
    - 0.9
    - 0.99
    weight_decay: 0.1
  scheduler:
    warmup_steps: 100
    constant_steps: 0
    decay_steps: 0
    min_lr: 1.0e-06
  model_ema_decay: 0.9999
world_model_metrics:
  _target_: mira.training.metrics.world_model_metrics.WorldModelMetricsConfig
  _convert_: all
  inference:
    schedule_type: linear
    noise_level: 0.0
  num_samples: 2048
  per_device_batch_size: 2
  n_context_frames: 38
  num_unrolled_frames: 20
  drift_metric_frames: 20
  fdd_slice_frames: 10
  eval_temporal_downsampling: null
  compile: false
  dino_max_chunk_size: 128
  num_viz_samples: 8
model:
  architecture:
    config:
      hidden_dim: 2048
      n_head: 16
      n_kv_head: 4
      n_layers: 16
      time_attention_every: 4
      _target_: mira.world_model.config.LatentWorldModelConfig
      actions:
        target_fps: 20
        valid_keys:
        - W
        - A
        - S
        - D
        - Q
        - E
        - Space
        - LShiftKey
        - LControlKey
      video:
        _target_: mira.ml.image_config.ImageConfig
        height: 288
        width: 512
        channels: 3
        timesteps: 80
        fps: 20
      codec_checkpoint: /weights/checkpoints/wm1b-psd10k-10000/codec/checkpoint-125000/checkpoint.pth
      latent_mean_std: null
      patch_size: 1
      n_register_tokens: 0
      attention_gating: true
      n_context_frames: 39
      dropout_action_prob: 0.1
      learned_temporal_pool: true
      causal: true
      use_codec_posterior_mean: true
      activation_checkpointing: psd-only
      ada_attn_ln: true
      psd_loss_prob: 0.1
      psd_weight: 0.0
      use_clean_past: true
    _target_: mira.world_model.latent_world_model.LatentWorldModel
    _convert_: all
dataset:
  actions:
    target_fps: 20
    valid_keys:
    - W
    - A
    - S
    - D
    - Q
    - E
    - Space
    - LShiftKey
    - LControlKey
  train_index: /home/hugohernandez/data/train
  test_index: /home/hugohernandez/data/test
  n_players: 1
  target_fps: 20
  frame_size: null
  exclude_replays: false