Refactor configs to have env in seperate yaml + Fix training

lerobot/configs/default.yaml

@@ -1,31 +1,24 @@
+defaults:
+  - _self_
+  - env: simxarm
+
 hydra:
   run:
     dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}_${hydra.job.name}
-  job:
-    name: default
 
 seed: 1337
 device: cuda
 buffer_device: cuda
-eval_freq: 1000
-save_freq: 10000
-eval_episodes: 20
+eval_freq: ???
+save_freq: ???
+eval_episodes: ???
 save_video: false
 save_model: false
 save_buffer: false
-train_steps: 50000
-fps: 15
-
-env:
-  name: simxarm
-  task: lift
-  from_pixels: True
-  pixels_only: False
-  image_size: 84
-  action_repeat: 2
-  episode_length: 25
-  fps: ${fps}
+train_steps: ???
+fps: ???
 
+env: ???
 
 policy:
   name: tdmpc
@@ -42,8 +35,8 @@ policy:
   frame_stack: 1
   num_channels: 32
   img_size: ${env.image_size}
-  state_dim: 4
-  action_dim: 4
+  state_dim: ???
+  action_dim: ???
 
   # planning
   mpc: true

lerobot/configs/env/pusht.yaml

@@ -1,6 +1,4 @@
-defaults:
-  - default
-  - _self_
+# @package _global_
 
 hydra:
   job:
@@ -9,11 +7,15 @@ hydra:
 eval_episodes: 50
 eval_freq: 7500
 save_freq: 75000
+train_steps: 50000  # TODO: same as simxarm, need to adjust
+
 fps: 10
 
 env:
   name: pusht
   task: pusht
+  from_pixels: True
+  pixels_only: False
   image_size: 96
   action_repeat: 1
   episode_length: 300

lerobot/configs/env/simxarm.yaml

@@ -0,0 +1,26 @@
+# @package _global_
+
+hydra:
+  job:
+    name: simxarm
+
+eval_episodes: 20
+eval_freq: 1000
+save_freq: 10000
+train_steps: 50000
+
+fps: 15
+
+env:
+  name: simxarm
+  task: lift
+  from_pixels: True
+  pixels_only: False
+  image_size: 84
+  action_repeat: 2
+  episode_length: 25
+  fps: ${fps}
+
+policy:
+  state_dim: 4
+  action_dim: 4