From 0eb9b5d1a54d618a023d75ec24499bb91b1cdb2c Mon Sep 17 00:00:00 2001
From: Cadene <re.cadene@gmail.com>
Date: Sun, 25 Feb 2024 11:15:09 +0000
Subject: [PATCH] Sanitize cfg.wandb

---
 lerobot/common/logger.py     | 11 +++++------
 lerobot/configs/default.yaml | 10 +++++-----
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/lerobot/common/logger.py b/lerobot/common/logger.py
index 1fd88714..306abd22 100644
--- a/lerobot/common/logger.py
+++ b/lerobot/common/logger.py
@@ -137,11 +137,10 @@ class Logger(object):
         self._cfg = cfg
         self._eval = []
         print_run(cfg)
-        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"
-        )
+        project = cfg.get("wandb", {}).get("project")
+        entity = cfg.get("wandb", {}).get("entity")
+        enable_wandb = cfg.get("wandb", {}).get("enable", False)
+        run_offline = not enable_wandb or not project or not entity
         if run_offline:
             print(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
             self._wandb = None
@@ -153,7 +152,7 @@ class Logger(object):
                 project=project,
                 entity=entity,
                 name=job_name,
-                notes=cfg.notes,
+                notes=cfg.get("wandb", {}).get("notes"),
                 # group=self._group,
                 tags=cfg_to_group(cfg, return_list=True),
                 dir=self._log_dir,
diff --git a/lerobot/configs/default.yaml b/lerobot/configs/default.yaml
index 922efac8..ad3b859e 100644
--- a/lerobot/configs/default.yaml
+++ b/lerobot/configs/default.yaml
@@ -107,8 +107,8 @@ policy:
   mlp_dim: 512
   latent_dim: 50
 
-# wandb
-use_wandb: true
-wandb_project: lerobot
-wandb_entity: rcadene  # insert your own
-notes: ""
+wandb:
+  enable: true
+  project: lerobot
+  entity: rcadene  # insert your own
+  notes: ""