Enable training with Tensorboard tracking

trlx/data/configs.py

@@ -137,8 +137,8 @@ class TrainConfig:
     :param batch_size: Batch size for training
     :type batch_size: int
 
-    :param trackers: Tuple of trackers to use for logging. Default: ("wandb",)
-    :type trackers: Tuple[str]
+    :param tracker: Tracker to use for logging. Default: "wandb"
+    :type tracker: str
 
     :param checkpoint_interval: Save model every checkpoint_interval steps
     :type checkpoint_interval: int
@@ -198,7 +198,9 @@ class TrainConfig:
     rollout_logging_dir: Optional[str] = None
     save_best: bool = True
 
-    trackers: Tuple[str] = ("wandb",)
+    tracker: Optional[str] = "wandb"
+    logging_dir: Optional[str] = None
+
     seed: int = 1000
 
     @classmethod

trlx/trainer/accelerate_base_trainer.py

@@ -44,7 +44,7 @@ class AccelerateRLTrainer(BaseRLTrainer):
     def __init__(self, config, **kwargs):
         super().__init__(config, **kwargs)
         self.max_length = config.train.seq_length
-        self.accelerator = Accelerator(log_with=config.train.trackers)
+        self.accelerator = Accelerator(log_with=config.train.tracker, logging_dir=config.train.logging_dir)
         if int(os.environ.get("WORLD_SIZE", 1)) > 1:
             torch.distributed.barrier(device_ids=[int(os.environ.get("LOCAL_RANK", 0))])
 
@@ -78,19 +78,32 @@ def __init__(self, config, **kwargs):
             dist_config = get_distributed_config(self.accelerator)
             config_dict["distributed"] = dist_config
             init_trackers_kwargs = {}
-            if "wandb" in config.train.trackers:
+            # HACK: Tensorboard doesn't like nested dict as hyperparams
+            config_dict_flat = {a:b for (k,v) in config_dict.items() for (a,b) in v.items() if not isinstance(b, dict)}
+
+            if config.train.tracker not in ("wandb", "tensorboard"):
+                raise ValueError(f"Only supported trackers are wandb and tensorboard, got {config.train.tracker}")
+
+            if config.train.tracker == "wandb":
                 init_trackers_kwargs["wandb"] = {
-                    "name": run_name,
-                    "entity": self.config.train.entity_name,
-                    "group": self.config.train.group_name,
-                    "tags": ["/".join(get_git_tag())],
-                    "mode": "disabled" if os.environ.get("debug", False) else "online",
-                }
-            self.accelerator.init_trackers(
-                project_name=self.config.train.project_name,
-                config=config_dict,
-                init_kwargs=init_trackers_kwargs,
-            )
+                        "name": run_name,
+                        "entity": self.config.train.entity_name,
+                        "group": self.config.train.group_name,
+                        "tags": ["/".join(get_git_tag())],
+                        "mode": "disabled" if os.environ.get("debug", False) else "online",
+                    }
+
+                self.accelerator.init_trackers(
+                    project_name=self.config.train.project_name,
+                    config=config_dict,
+                    init_kwargs=init_trackers_kwargs,
+                )
+            else:
+                self.accelerator.init_trackers(
+                    project_name=self.config.train.project_name,
+                    config=config_dict_flat,
+                )
+
 
     def setup_model(self):
         """