mira-mini-psd / world_model_config.yaml
Karajan42's picture
weights: 1B PSD 2-step bundle (wm1b-psd10k-10000, verified load + real context)
ec6b120 verified
Raw
History Blame Contribute Delete
2.8 kB
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