Add training config

config.yaml

@@ -0,0 +1,71 @@
+output_model_name: lewm
+subdir: ${hydra:job.id}
+num_workers: 1
+train_split: 0.9
+seed: 3072
+img_size: 224
+patch_size: 14
+encoder_scale: tiny
+encoder_resnet9: false
+train_value_function: false
+projector_loss_weight: 0.001
+dump_object: true
+trainer:
+  max_epochs: 5
+  devices: auto
+  accelerator: gpu
+  precision: bf16
+  gradient_clip_val: 1.0
+loader:
+  batch_size: 128
+  num_workers: 1
+  drop_last: true
+  persistent_workers: false
+  prefetch_factor: 3
+  pin_memory: true
+  shuffle: true
+optimizer:
+  type: AdamW
+  lr: 5.0e-05
+  weight_decay: 0.001
+wandb:
+  enabled: true
+  config:
+    entity: paoloai-robomotic
+    project: lewm
+    name: ${output_model_name}
+    id: ${subdir}
+    resume: allow
+    log_model: false
+wm:
+  type: lewm
+  history_size: 3
+  num_preds: 1
+  embed_dim: 192
+  action_dim: 2
+  proprio_dim: 2
+predictor:
+  depth: 6
+  heads: 16
+  mlp_dim: 2048
+  dim_head: 64
+  dropout: 0.1
+  emb_dropout: 0.0
+loss:
+  sigreg:
+    weight: 0.09
+    kwargs:
+      knots: 17
+      num_proj: 1024
+data:
+  dataset:
+    num_steps: ${eval:'${wm.num_preds} + ${wm.history_size}'}
+    frameskip: 5
+    name: glitched_hue_tworoom
+    keys_to_load:
+    - pixels
+    - action
+    - proprio
+    keys_to_cache:
+    - action
+    - proprio