[Feature] Expose additional arguments for Wandb

benchmarl/conf/experiment/base_experiment.yaml

@@ -101,8 +101,12 @@ evaluation_static: False
 
 # List of loggers to use, options are: wandb, csv, tensorboard, mflow
 loggers: [csv,wandb]
-# Wandb project name
+# Wandb project name (kept for backward compatibility)
 project_name: "benchmarl"
+# Wandb extra kwargs passed to the WandbLogger (~superset of wandb.init kwargs)
+# WandbLogger includes: offline, save_dir, project, video_fps
+# wandb.init includes: entity, tags, notes, etc.
+wandb_extra_kwargs: {}
 # Create a json folder as part of the output in the format of marl-eval
 create_json: True
 

benchmarl/experiment/experiment.py

@@ -111,6 +111,7 @@ class ExperimentConfig:
 
     loggers: List[str] = MISSING
     project_name: str = MISSING
+    wandb_extra_kwargs: Dict[str, Any] = MISSING
     create_json: bool = MISSING
 
     save_folder: Optional[str] = MISSING
@@ -611,7 +612,6 @@ def _setup_name(self):
 
     def _setup_logger(self):
         self.logger = Logger(
-            project_name=self.config.project_name,
             experiment_name=self.name,
             folder_name=str(self.folder_name),
             experiment_config=self.config,
@@ -621,6 +621,8 @@ def _setup_logger(self):
             task_name=self.task_name,
             group_map=self.group_map,
             seed=self.seed,
+            project_name=self.config.project_name,
+            wandb_extra_kwargs=self.config.wandb_extra_kwargs,
         )
         self.logger.log_hparams(
             critic_model_name=self.critic_model_name,

benchmarl/experiment/logger.py

@@ -10,7 +10,7 @@
 from collections.abc import MutableMapping, Sequence
 from pathlib import Path
 
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 import numpy as np
 import torch
@@ -39,6 +39,7 @@ def __init__(
         group_map: Dict[str, List[str]],
         seed: int,
         project_name: str,
+        wandb_extra_kwargs: Dict[str, Any],
     ):
         self.experiment_config = experiment_config
         self.algorithm_name = algorithm_name
@@ -62,15 +63,21 @@ def __init__(
 
         self.loggers: List[torchrl.record.loggers.Logger] = []
         for logger_name in experiment_config.loggers:
+            wandb_project = wandb_extra_kwargs.get("project", project_name)
+            if wandb_project != project_name:
+                raise ValueError(
+                    f"wandb_extra_kwargs.project ({wandb_project}) is different from the project_name ({project_name})"
+                )
             self.loggers.append(
                 get_logger(
                     logger_type=logger_name,
                     logger_name=folder_name,
                     experiment_name=experiment_name,
                     wandb_kwargs={
                         "group": task_name,
-                        "project": project_name,
                         "id": experiment_name,
+                        "project": project_name,
+                        **wandb_extra_kwargs,
                     },
                 )
             )

benchmarl/hydra_config.py

@@ -75,9 +75,10 @@ def load_task_config_from_hydra(cfg: DictConfig, task_name: str) -> TaskClass:
     cfg_dict_checked = OmegaConf.to_object(cfg)
     if is_dataclass(cfg_dict_checked):
         cfg_dict_checked = cfg_dict_checked.__dict__
-    cfg_dict_checked = _type_check_task_config(
-        environment_name, inner_task_name, cfg_dict_checked
-    )  # Only needed for the warning
+    else:
+        cfg_dict_checked = _type_check_task_config(
+            environment_name, inner_task_name, cfg_dict_checked
+        )  # Only needed for the warning
     return task_config_registry[task_name].get_task(cfg_dict_checked)