Changes Wandb runner name to log directory name

Approved-by: Clemens Schwarke

Author: Mayankm96

rsl_rl/utils/wandb_utils.py

@@ -21,6 +21,9 @@ class WandbSummaryWriter(SummaryWriter):
     def __init__(self, log_dir: str, flush_secs: int, cfg):
         super().__init__(log_dir, flush_secs)
 
+        # Get the run name
+        run_name = os.path.split(log_dir)[-1]
+
         try:
             project = cfg["wandb_project"]
         except KeyError:
@@ -29,35 +32,27 @@ def __init__(self, log_dir: str, flush_secs: int, cfg):
         try:
             entity = os.environ["WANDB_USERNAME"]
         except KeyError:
-            raise KeyError(
-                "Wandb username not found. Please run or add to ~/.bashrc: export WANDB_USERNAME=YOUR_USERNAME"
-            )
+            entity = None
 
-        wandb.init(project=project, entity=entity)
+        # Initialize wandb
+        wandb.init(project=project, entity=entity, name=run_name)
 
-        # Change generated name to project-number format
-        wandb.run.name = project + wandb.run.name.split("-")[-1]
+        # Add log directory to wandb
+        wandb.config.update({"log_dir": log_dir})
 
         self.name_map = {
             "Train/mean_reward/time": "Train/mean_reward_time",
             "Train/mean_episode_length/time": "Train/mean_episode_length_time",
         }
 
-        run_name = os.path.split(log_dir)[-1]
-
-        wandb.log({"log_dir": run_name})
-
     def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg):
         wandb.config.update({"runner_cfg": runner_cfg})
         wandb.config.update({"policy_cfg": policy_cfg})
         wandb.config.update({"alg_cfg": alg_cfg})
-        wandb.config.update({"env_cfg": asdict(env_cfg)})
-
-    def _map_path(self, path):
-        if path in self.name_map:
-            return self.name_map[path]
-        else:
-            return path
+        try:
+            wandb.config.update({"env_cfg": env_cfg.to_dict()})
+        except Exception:
+            wandb.config.update({"env_cfg": asdict(env_cfg)})
 
     def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False):
         super().add_scalar(
@@ -80,3 +75,13 @@ def save_model(self, model_path, iter):
 
     def save_file(self, path, iter=None):
         wandb.save(path, base_path=os.path.dirname(path))
+
+    """
+    Private methods.
+    """
+
+    def _map_path(self, path):
+        if path in self.name_map:
+            return self.name_map[path]
+        else:
+            return path