Try to make tensorboard properly optional

coltra/policy_optimization.py

@@ -8,7 +8,6 @@
 from numpy.random import Generator
 import torch
 from torch import Tensor
-from torch.utils.tensorboard import SummaryWriter
 from typarse import BaseConfig
 
 from coltra.agents import Agent
@@ -114,7 +113,7 @@ def train_on_data(
         data_dict: dict[str, OnPolicyRecord],
         shape: tuple[int, int],
         step: int = 0,
-        writer: Optional[SummaryWriter] = None,
+        writer: Optional["SummaryWriter"] = None,
     ) -> dict[str, float]:
         """
         Performs a single update step with PPO on the given batch of data.

coltra/research/policy_fusion/__init__.py

@@ -1 +0,0 @@
-from .policy_fusion import JointModel

coltra/trainers.py

@@ -9,7 +9,6 @@
 import shortuuid
 import torch
 import yaml
-from torch.utils.tensorboard import SummaryWriter
 from tqdm import trange
 
 from coltra.collectors import collect_crowd_data, collect_family_data
@@ -67,8 +66,10 @@ def __init__(
         )
 
         # Setup tensorboard
-        self.writer: Optional[SummaryWriter]
+        # self.writer: Optional[SummaryWriter]
         if self.config.tensorboard_name:
+            from torch.utils.tensorboard import SummaryWriter
+
             dt_string = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
             if use_uuid:
                 dt_string += "_" + shortuuid.uuid()
@@ -242,8 +243,9 @@ def __init__(
         )
 
         # Setup tensorboard
-        self.writer: Optional[SummaryWriter]
         if self.config.tensorboard_name:
+            from torch.utils.tensorboard import SummaryWriter
+
             dt_string = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
             if use_uuid:
                 dt_string += "_" + shortuuid.uuid()

coltra/utils.py

@@ -25,7 +25,7 @@
 from torch.optim.adamax import Adamax
 from torch.optim.sgd import SGD
 
-from torch.utils.tensorboard import SummaryWriter
+# from torch.utils.tensorboard import SummaryWriter
 
 from coltra.buffers import Observation, Action
 
@@ -37,7 +37,7 @@
 def write_dict(
     metrics: dict[str, Union[int, float]],
     step: int,
-    writer: Optional[SummaryWriter] = None,
+    writer: Optional["SummaryWriter"] = None,
 ):
     """Writes a dictionary to a tensorboard SummaryWriter"""
     if writer is not None: