Wandb works, One output dir

lerobot/common/envs/factory.py

@@ -11,6 +11,7 @@ def make_env(cfg):
         "from_pixels": cfg.from_pixels,
         "pixels_only": cfg.pixels_only,
         "image_size": cfg.image_size,
+        "max_episode_length": cfg.episode_length,
     }
 
     if cfg.env == "simxarm":

lerobot/common/envs/pusht.py

@@ -29,7 +29,7 @@ class PushtEnv(EnvBase):
         image_size=None,
         seed=1337,
         device="cpu",
-        max_episode_length=25,  # TODO: verify
+        max_episode_length=300,
     ):
         super().__init__(device=device, batch_size=[])
         self.frame_skip = frame_skip
@@ -53,13 +53,11 @@ class PushtEnv(EnvBase):
         if not from_pixels:
             raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv")
         from diffusion_policy.env.pusht.pusht_image_env import PushTImageEnv
-        from gym.wrappers import TimeLimit
 
         self._env = PushTImageEnv(render_size=self.image_size)
-        self._env = TimeLimit(self._env, self.max_episode_length)
 
         self._make_spec()
-        self.set_seed(seed)
+        self._current_seed = self.set_seed(seed)
 
     def render(self, mode="rgb_array", width=384, height=384):
         if width != height:
@@ -90,7 +88,11 @@ class PushtEnv(EnvBase):
     def _reset(self, tensordict: Optional[TensorDict] = None):
         td = tensordict
         if td is None or td.is_empty():
+            # we need to handle seed iteration, since self._env.reset() rely an internal _seed.
+            self._current_seed += 1
+            self.set_seed(self._current_seed)
             raw_obs = self._env.reset()
+            assert self._current_seed == self._env._seed
 
             td = TensorDict(
                 {

lerobot/common/envs/simxarm.py

@@ -49,7 +49,6 @@ class SimxarmEnv(EnvBase):
             raise ImportError("Cannot import gym.")
 
         import gym
-        from gym.wrappers import TimeLimit
         from simxarm import TASKS
 
         if self.task not in TASKS:
@@ -58,7 +57,6 @@ class SimxarmEnv(EnvBase):
             )
 
         self._env = TASKS[self.task]["env"]()
-        self._env = TimeLimit(self._env, TASKS[self.task]["episode_length"])
 
         MAX_NUM_ACTIONS = 4
         num_actions = len(TASKS[self.task]["action_space"])

lerobot/common/logger.py

@@ -11,8 +11,9 @@ from termcolor import colored
 CONSOLE_FORMAT = [
     ("episode", "E", "int"),
     ("env_step", "S", "int"),
-    ("avg_reward", "R", "float"),
-    ("pc_success", "R", "float"),
+    ("avg_sum_reward", "RS", "float"),
+    ("avg_max_reward", "RM", "float"),
+    ("pc_success", "S", "float"),
     ("total_time", "T", "time"),
 ]
 AGENT_METRICS = [
@@ -69,7 +70,11 @@ def print_run(cfg, reward=None):
 
 def cfg_to_group(cfg, return_list=False):
     """Return a wandb-safe group name for logging. Optionally returns group name as list."""
-    lst = [cfg.task, cfg.modality, re.sub("[^0-9a-zA-Z]+", "-", cfg.exp_name)]
+    # lst = [cfg.task, cfg.modality, re.sub("[^0-9a-zA-Z]+", "-", cfg.exp_name)]
+    lst = [
+        f"env:{cfg.env}",
+        f"seed:{cfg.seed}",
+    ]
     return lst if return_list else "-".join(lst)
 
 
@@ -120,8 +125,9 @@ class VideoRecorder:
 class Logger(object):
     """Primary logger object. Logs either locally or using wandb."""
 
-    def __init__(self, log_dir, cfg):
+    def __init__(self, log_dir, job_name, cfg):
         self._log_dir = make_dir(Path(log_dir))
+        self._job_name = job_name
         self._model_dir = make_dir(self._log_dir / "models")
         self._buffer_dir = make_dir(self._log_dir / "buffers")
         self._save_model = cfg.save_model
@@ -131,9 +137,8 @@ class Logger(object):
         self._cfg = cfg
         self._eval = []
         print_run(cfg)
-        project, entity = cfg.get("wandb_project", "none"), cfg.get(
-            "wandb_entity", "none"
-        )
+        project = cfg.get("wandb_project", "none")
+        entity = cfg.get("wandb_entity", "none")
         run_offline = (
             not cfg.get("use_wandb", False) or project == "none" or entity == "none"
         )
@@ -141,35 +146,39 @@ class Logger(object):
             print(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
             self._wandb = None
         else:
-            try:
-                os.environ["WANDB_SILENT"] = "true"
-                import wandb
+            # try:
+            os.environ["WANDB_SILENT"] = "true"
+            import wandb
 
-                wandb.init(
-                    project=project,
-                    entity=entity,
-                    name=str(cfg.seed),
-                    notes=cfg.notes,
-                    group=self._group,
-                    tags=cfg_to_group(cfg, return_list=True) + [f"seed:{cfg.seed}"],
-                    dir=self._log_dir,
-                    config=OmegaConf.to_container(cfg, resolve=True),
-                )
-                print(
-                    colored("Logs will be synced with wandb.", "blue", attrs=["bold"])
-                )
-                self._wandb = wandb
-            except:
-                print(
-                    colored(
-                        "Warning: failed to init wandb. Make sure `wandb_entity` is set to your username in `config.yaml`. Logs will be saved locally.",
-                        "yellow",
-                        attrs=["bold"],
-                    )
-                )
-                self._wandb = None
+            wandb.init(
+                project=project,
+                entity=entity,
+                name=job_name,
+                notes=cfg.notes,
+                # group=self._group,
+                tags=cfg_to_group(cfg, return_list=True),
+                dir=self._log_dir,
+                config=OmegaConf.to_container(cfg, resolve=True),
+                # TODO(rcadene): try set to True
+                save_code=False,
+                # TODO(rcadene): split train and eval, and run async eval with job_type="eval"
+                job_type="train_eval",
+                # TODO(rcadene): add resume option
+                resume=None,
+            )
+            print(colored("Logs will be synced with wandb.", "blue", attrs=["bold"]))
+            self._wandb = wandb
+            # except:
+            #     print(
+            #         colored(
+            #             "Warning: failed to init wandb. Make sure `wandb_entity` is set to your username in `config.yaml`. Logs will be saved locally.",
+            #             "yellow",
+            #             attrs=["bold"],
+            #         )
+            #     )
+            #     self._wandb = None
         self._video = (
-            VideoRecorder(log_dir, self._wandb)
+            VideoRecorder(self._log_dir, self._wandb)
             if self._wandb and cfg.save_video
             else None
         )
@@ -235,7 +244,7 @@ class Logger(object):
                 self._wandb.log({category + "/" + k: v}, step=d["env_step"])
         if category == "eval":
             # keys = ['env_step', 'avg_reward']
-            keys = ["env_step", "avg_reward", "pc_success"]
+            keys = ["env_step", "avg_sum_reward", "avg_max_reward", "pc_success"]
             self._eval.append(np.array([d[key] for key in keys]))
             pd.DataFrame(np.array(self._eval)).to_csv(
                 self._log_dir / "eval.log", header=keys, index=None

lerobot/common/tdmpc.py

@@ -96,7 +96,7 @@ class TDMPC(nn.Module):
         self.model_target.eval()
         self.batch_size = cfg.batch_size
 
-        self.step = 0
+        self.register_buffer("step", torch.zeros(1))
 
     def state_dict(self):
         """Retrieve state dict of TOLD model, including slow-moving target network."""
@@ -122,7 +122,7 @@ class TDMPC(nn.Module):
             "rgb": observation["image"],
             "state": observation["state"],
         }
-        return self.act(obs, t0=t0, step=self.step)
+        return self.act(obs, t0=t0, step=self.step.item())
 
     @torch.no_grad()
     def act(self, obs, t0=False, step=None):
@@ -513,5 +513,5 @@ class TDMPC(nn.Module):
         metrics.update(value_info)
         metrics.update(pi_update_info)
 
-        self.step = step
+        self.step[0] = step
         return metrics