output_model_name: lewm
subdir: ${hydra:job.id}
save_dir: xxx
num_workers: 8
train_split: 0.9
seed: 3072
img_size: 224
patch_size: 14
encoder_scale: tiny
dump_object: true
trainer:
  max_epochs: 10
  devices: 1
  accelerator: gpu
  precision: bf16
  gradient_clip_val: 1.0
loader:
  batch_size: 128
  num_workers: ${num_workers}
  persistent_workers: true
  prefetch_factor: 3
  pin_memory: true
wm:
  type: lewm
  history_size: 3
  num_preds: 1
  embed_dim: 192
  sigreg_subspaces: 32
  sigreg_subspace_dim: 6
  sigreg_init_mode: 2
  action_dim: 5
  observation_dim: 28
  proprio_dim: 19
predictor:
  depth: 6
  heads: 16
  mlp_dim: 2048
  dim_head: 64
  dropout: 0.1
  emb_dropout: 0.0
loss:
  sigreg:
    weight: 0.09
    theta: 0.0
    kwargs:
      knots: 17
      num_proj: 1024
optimizer:
  type: AdamW
  lr: 5.0e-05
  weight_decay: 0.001
wandb:
  enabled: true
  config:
    entity: xxx
    project: lewm
    name: ${output_model_name}
    id: ${subdir}
    resume: allow
    log_model: false
data:
  dataset:
    name: cube_single_expert
    num_steps: ${eval:'${wm.num_preds} + ${wm.history_size}'}
    frameskip: 5
    keys_to_load:
    - pixels
    - action
    - observation
    - proprio
    keys_to_cache:
    - action
    - observation
    keys_to_merge:
      proprio: proprio