Sanitize cfg.env

lerobot/common/envs/factory.py

@@ -7,26 +7,26 @@ from lerobot.common.envs.transforms import Prod
 
 def make_env(cfg):
     kwargs = {
-        "frame_skip": cfg.action_repeat,
-        "from_pixels": cfg.from_pixels,
-        "pixels_only": cfg.pixels_only,
-        "image_size": cfg.image_size,
+        "frame_skip": cfg.env.action_repeat,
+        "from_pixels": cfg.env.from_pixels,
+        "pixels_only": cfg.env.pixels_only,
+        "image_size": cfg.env.image_size,
     }
 
-    if cfg.env == "simxarm":
-        kwargs["task"] = cfg.task
+    if cfg.env.name == "simxarm":
+        kwargs["task"] = cfg.env.task
         clsfunc = SimxarmEnv
-    elif cfg.env == "pusht":
+    elif cfg.env.name == "pusht":
         clsfunc = PushtEnv
     else:
-        raise ValueError(cfg.env)
+        raise ValueError(cfg.env.name)
 
     env = clsfunc(**kwargs)
 
     # limit rollout to max_steps
-    env = TransformedEnv(env, StepCounter(max_steps=cfg.episode_length))
+    env = TransformedEnv(env, StepCounter(max_steps=cfg.env.episode_length))
 
-    if cfg.env == "pusht":
+    if cfg.env.name == "pusht":
         # to ensure pusht is in [0,255] like simxarm
         env.append_transform(Prod(in_keys=[("observation", "image")], prod=255.0))
 

lerobot/common/logger.py

@@ -51,7 +51,7 @@ def print_run(cfg, reward=None):
 
     kvs = [
         ("task", cfg.task),
-        ("train steps", f"{int(cfg.train_steps * cfg.action_repeat):,}"),
+        ("train steps", f"{int(cfg.train_steps * cfg.env.action_repeat):,}"),
         # ('observations', 'x'.join([str(s) for s in cfg.obs_shape])),
         # ('actions', cfg.action_dim),
         # ('experiment', cfg.exp_name),
@@ -117,7 +117,11 @@ class VideoRecorder:
         if self.enabled:
             frames = np.stack(self.frames).transpose(0, 3, 1, 2)
             self._wandb.log(
-                {"eval_video": self._wandb.Video(frames, fps=self.fps, format="mp4")},
+                {
+                    "eval_video": self._wandb.Video(
+                        frames, fps=self.env.fps, format="mp4"
+                    )
+                },
                 step=step,
             )