Tensorboard integration (#30)

* init commit tensorboard-integration

* Added tb logger to ppo (with output exclusions)

* fixed truncated stdout

* categorize stdout outputs by tag

* separated exclusions from values, added missing logs

* saving exclusions as dict instead of list

* reformatting, auto run indexing

* included renaming suggestions, fixed tests

* tb support for sac

* linting

* moved logging to base class

* tb support for td3

* removed histograms, non-verbose output working

* modifed changelog

* linting

* fixed type error

* moved logger config to utils

* removed episode_rewards log from ppo

* Enable tensorboard in tests

* Remove unused import

* Update logger sub titles

* Minor edit for PPO

* Update logger and tb log folder

* Pass correct logger to Callbacks

* updated docs

* added tb example image to docs

* add support for continuing training in tensorboard

* added tensorboard to docs index

* added tb test

* moved logger config to _setup_learn, updated tests

* accessing verbose from base class

* Update doc and tests

* Rename session -> time

* Update version

* Update logger truncate

* Update types

* Remove duplicated code

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

README.md

@@ -29,6 +29,7 @@ These algorithms will make it easier for the research community and industry to
 | Custom policies             | :heavy_check_mark: |
 | Common interface            | :heavy_check_mark: |
 | Ipython / Notebook friendly | :heavy_check_mark: |
+| Tensorboard support         | :heavy_check_mark: |
 | PEP8 code style             | :heavy_check_mark: |
 | Custom callback             | :heavy_check_mark: |
 | High code coverage          | :heavy_check_mark: |
@@ -48,7 +49,6 @@ Planned features:
 
 ### Planned features (v1.1+)
 
-- [ ] Full Tensorboard support
 - [ ] DQN extensions (prioritized replay, double q-learning, ...)
 - [ ] Support for `Tuple` and `Dict` observation spaces
 - [ ] Recurrent Policies
@@ -99,7 +99,7 @@ Install the Stable Baselines3 package:
 pip install stable-baselines3[extra]
 ```
 
-This includes an optional dependencies like OpenCV or `atari-py` to train on atari games. If you do not need those, you can use:
+This includes an optional dependencies like Tensorboard, OpenCV or `atari-py` to train on atari games. If you do not need those, you can use:
 ```
 pip install stable-baselines3
 ```

docs/guide/callbacks.rst

@@ -44,7 +44,7 @@ This will give you access to events (``_on_training_start``, ``_on_step``) and u
             # self.locals = None  # type: Dict[str, Any]
             # self.globals = None  # type: Dict[str, Any]
             # The logger object, used to report things in the terminal
-            # self.logger = None  # type: logger.Logger
+            # self.logger = None  # stable_baselines3.common.logger
             # # Sometimes, for event callback, it is useful
             # # to have access to the parent object
             # self.parent = None  # type: Optional[BaseCallback]

docs/guide/install.rst

@@ -29,7 +29,7 @@ To install Stable Baselines3 with pip, execute:
 
     pip install stable-baselines3[extra]
 
-This includes an optional dependencies like OpenCV or ```atari-py``` to train on atari games. If you do not need those, you can use:
+This includes an optional dependencies like Tensorboard, OpenCV or ```atari-py``` to train on atari games. If you do not need those, you can use:
 
 .. code-block:: bash
 

docs/guide/tensorboard.rst

@@ -0,0 +1,82 @@
+.. _tensorboard:
+
+Tensorboard Integration
+=======================
+
+Basic Usage
+------------
+
+To use Tensorboard with stable baselines3, you simply need to pass the location of the log folder to the RL agent:
+
+.. code-block:: python
+
+    from stable_baselines3 import A2C
+
+    model = A2C('MlpPolicy', 'CartPole-v1', verbose=1, tensorboard_log="./a2c_cartpole_tensorboard/")
+    model.learn(total_timesteps=10000)
+
+
+You can also define custom logging name when training (by default it is the algorithm name)
+
+.. code-block:: python
+
+    from stable_baselines3 import A2C
+
+    model = A2C('MlpPolicy', 'CartPole-v1', verbose=1, tensorboard_log="./a2c_cartpole_tensorboard/")
+    model.learn(total_timesteps=10000, tb_log_name="first_run")
+    # Pass reset_num_timesteps=False to continue the training curve in tensorboard
+    # By default, it will create a new curve
+    model.learn(total_timesteps=10000, tb_log_name="second_run", reset_num_timesteps=False)
+    model.learn(total_timesteps=10000, tb_log_name="third_run", reset_num_timesteps=False)
+
+
+Once the learn function is called, you can monitor the RL agent during or after the training, with the following bash command:
+
+.. code-block:: bash
+
+  tensorboard --logdir ./a2c_cartpole_tensorboard/
+
+you can also add past logging folders:
+
+.. code-block:: bash
+
+  tensorboard --logdir ./a2c_cartpole_tensorboard/;./ppo2_cartpole_tensorboard/
+
+It will display information such as the episode reward (when using a ``Monitor`` wrapper), the model losses and other parameter unique to some models.
+
+.. image:: ../_static/img/Tensorboard_example.png
+  :width: 600
+  :alt: plotting
+
+Logging More Values
+-------------------
+
+Using a callback, you can easily log more values with TensorBoard.
+Here is a simple example on how to log both additional tensor or arbitrary scalar value:
+
+.. code-block:: python
+
+    import numpy as np
+
+    from stable_baselines3 import SAC
+    from stable_baselines3.common.callbacks import BaseCallback
+
+    model = SAC("MlpPolicy", "Pendulum-v0", tensorboard_log="/tmp/sac/", verbose=1)
+
+
+    class TensorboardCallback(BaseCallback):
+        """
+        Custom callback for plotting additional values in tensorboard.
+        """
+
+        def __init__(self, verbose=0):
+            super(TensorboardCallback, self).__init__(verbose)
+
+        def _on_step(self) -> bool:
+            # Log scalar value (here a random variable)
+            value = np.random.random()
+            self.logger.record('random_value', value)
+            return True
+
+
+    model.learn(50000, callback=TensorboardCallback())

docs/index.rst

@@ -25,7 +25,7 @@ Main Features
 - Documented functions and classes
 - Tests, high code coverage and type hints
 - Clean code
-
+- Tensorboard support
 
 
 .. toctree::
@@ -42,6 +42,7 @@ Main Features
    guide/custom_env
    guide/custom_policy
    guide/callbacks
+   guide/tensorboard
    guide/rl_zoo
    guide/migration
    guide/checking_nan

docs/misc/changelog.rst

@@ -3,12 +3,17 @@
 Changelog
 ==========
 
-Pre-Release 0.6.0a10 (WIP)
+Pre-Release 0.6.0a11 (WIP)
 ------------------------------
 
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - Remove State-Dependent Exploration (SDE) support for ``TD3``
+- Methods were renamed in the logger:
+  - ``logkv`` -> ``record``, ``writekvs`` -> ``write``, ``writeseq`` ->  ``write_sequence``,
+  - ``logkvs`` -> ``record_dict``, ``dumpkvs`` -> ``dump``,
+  - ``getkvs`` -> ``get_log_dict``, ``logkv_mean`` -> ``record_mean``,
+
 
 New Features:
 ^^^^^^^^^^^^^
@@ -18,7 +23,9 @@ New Features:
 - Added ``cmd_util`` and ``atari_wrappers``
 - Added support for ``MultiDiscrete`` and ``MultiBinary`` observation spaces (@rolandgvc)
 - Added ``MultiCategorical`` and ``Bernoulli`` distributions for PPO/A2C (@rolandgvc)
+- Added support for logging to tensorboard (@rolandgvc)
 - Added ``VectorizedActionNoise`` for continuous vectorized environments (@PartiallyTyped)
+- Log evaluation in the ``EvalCallback`` using the logger
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -26,6 +33,7 @@ Bug Fixes:
 - Fixed version number that had a new line included
 - Fixed weird seg fault in docker image due to FakeImageEnv by reducing screen size
 - Fixed ``sde_sample_freq`` that was not taken into account for SAC
+- Pass logger module to ``BaseCallback`` otherwise they cannot write in the one used by the algorithms
 
 Deprecations:
 ^^^^^^^^^^^^^

setup.py

@@ -106,7 +106,9 @@ setup(name='stable_baselines3',
               # For render
               'opencv-python',
               # For atari games,
-              'atari_py~=0.2.0', 'pillow'
+              'atari_py~=0.2.0', 'pillow',
+              # Tensorboard support
+              'tensorboard'
           ]
       },
       description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',

stable_baselines3/a2c/a2c.py

@@ -141,13 +141,13 @@ class A2C(PPO):
                                            self.rollout_buffer.values.flatten())
 
         self._n_updates += 1
-        logger.logkv("n_updates", self._n_updates)
+        logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
-        logger.logkv("explained_variance", explained_var)
+        logger.record("train/explained_variance", explained_var)
-        logger.logkv("entropy_loss", entropy_loss.item())
+        logger.record("train/entropy_loss", entropy_loss.item())
-        logger.logkv("policy_loss", policy_loss.item())
+        logger.record("train/policy_loss", policy_loss.item())
-        logger.logkv("value_loss", value_loss.item())
+        logger.record("train/value_loss", value_loss.item())
         if hasattr(self.policy, 'log_std'):
-            logger.logkv("std", th.exp(self.policy.log_std).mean().item())
+            logger.record("train/std", th.exp(self.policy.log_std).mean().item())
 
     def learn(self,
               total_timesteps: int,

stable_baselines3/common/base_class.py

@@ -11,7 +11,7 @@ import gym
 import torch as th
 import numpy as np
 
-from stable_baselines3.common import logger
+from stable_baselines3.common import logger, utils
 from stable_baselines3.common.policies import BasePolicy, get_policy_from_name
 from stable_baselines3.common.utils import set_random_seed, get_schedule_fn, update_learning_rate, get_device
 from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize, VecTransposeImage
@@ -35,6 +35,7 @@ class BaseRLModel(ABC):
     :param learning_rate: (float or callable) learning rate for the optimizer,
         it can be a function of the current progress (from 1 to 0)
     :param policy_kwargs: (Dict[str, Any]) Additional arguments to be passed to the policy on creation
+    :param tensorboard_log: (str) the log location for tensorboard (if None, no logging)
     :param verbose: (int) The verbosity level: 0 none, 1 training information, 2 debug
     :param device: (Union[th.device, str]) Device on which the code should run.
         By default, it will try to use a Cuda compatible device and fallback to cpu
@@ -58,6 +59,7 @@ class BaseRLModel(ABC):
                  policy_base: Type[BasePolicy],
                  learning_rate: Union[float, Callable],
                  policy_kwargs: Dict[str, Any] = None,
+                 tensorboard_log: Optional[str] = None,
                  verbose: int = 0,
                  device: Union[th.device, str] = 'auto',
                  support_multi_env: bool = False,
@@ -91,6 +93,7 @@ class BaseRLModel(ABC):
         self.start_time = None
         self.policy = None
         self.learning_rate = learning_rate
+        self.tensorboard_log = tensorboard_log
         self.lr_schedule = None  # type: Optional[Callable]
         self._last_obs = None  # type: Optional[np.ndarray]
         # When using VecNormalize:
@@ -191,7 +194,7 @@ class BaseRLModel(ABC):
             An optimizer or a list of optimizers.
         """
         # Log the current learning rate
-        logger.logkv("learning_rate", self.lr_schedule(self._current_progress))
+        logger.record("train/learning_rate", self.lr_schedule(self._current_progress))
 
         if not isinstance(optimizers, list):
             optimizers = [optimizers]
@@ -289,7 +292,7 @@ class BaseRLModel(ABC):
         """
         Return a trained model.
 
-        :param total_timesteps: (int) The total number of samples to train on
+        :param total_timesteps: (int) The total number of samples (env steps) to train on
         :param callback: (function (dict, dict)) -> boolean function called at every steps with state of the algorithm.
             It takes the local and global variables. If it returns False, training is aborted.
         :param log_interval: (int) The number of timesteps before logging.
@@ -491,23 +494,27 @@ class BaseRLModel(ABC):
         return callback
 
     def _setup_learn(self,
+                     total_timesteps: int,
                      eval_env: Optional[GymEnv],
                      callback: Union[None, Callable, List[BaseCallback], BaseCallback] = None,
                      eval_freq: int = 10000,
                      n_eval_episodes: int = 5,
                      log_path: Optional[str] = None,
                      reset_num_timesteps: bool = True,
-                     ) -> 'BaseCallback':
+                     tb_log_name: str = 'run',
+                     ) -> Tuple[int, 'BaseCallback']:
         """
         Initialize different variables needed for training.
 
+        :param total_timesteps: (int) The total number of samples (env steps) to train on
         :param eval_env: (Optional[GymEnv])
         :param callback: (Union[None, BaseCallback, List[BaseCallback, Callable]])
         :param eval_freq: (int)
         :param n_eval_episodes: (int)
         :param log_path (Optional[str]): Path to a log folder
         :param reset_num_timesteps: (bool) Whether to reset or not the ``num_timesteps`` attribute
-        :return: (BaseCallback)
+        :param tb_log_name: (str) the name of the run for tensorboard log
+        :return: (int, Tuple[BaseCallback])
         """
         self.start_time = time.time()
         self.ep_info_buffer = deque(maxlen=100)
@@ -519,6 +526,9 @@ class BaseRLModel(ABC):
         if reset_num_timesteps:
             self.num_timesteps = 0
             self._episode_num = 0
+        else:
+            # Make sure training timesteps are ahead of the internal counter
+            total_timesteps += self.num_timesteps
 
         # Avoid resetting the environment when calling ``.learn()`` consecutive times
         if reset_num_timesteps or self._last_obs is None:
@@ -532,10 +542,13 @@ class BaseRLModel(ABC):
 
         eval_env = self._get_eval_env(eval_env)
 
+        # Configure logger's outputs
+        utils.configure_logger(self.verbose, self.tensorboard_log, tb_log_name, reset_num_timesteps)
+
         # Create eval callback if needed
         callback = self._init_callback(callback, eval_env, eval_freq, n_eval_episodes, log_path)
 
-        return callback
+        return total_timesteps, callback
 
     def _update_info_buffer(self, infos: List[Dict[str, Any]], dones: Optional[np.ndarray] = None) -> None:
         """
@@ -697,6 +710,7 @@ class OffPolicyRLModel(BaseRLModel):
                  learning_starts: int = 100,
                  batch_size: int = 256,
                  policy_kwargs: Dict[str, Any] = None,
+                 tensorboard_log: Optional[str] = None,
                  verbose: int = 0,
                  device: Union[th.device, str] = 'auto',
                  support_multi_env: bool = False,
@@ -709,13 +723,13 @@ class OffPolicyRLModel(BaseRLModel):
                  sde_support: bool = True):
 
         super(OffPolicyRLModel, self).__init__(policy, env, policy_base, learning_rate,
-                                               policy_kwargs, verbose,
+                                               policy_kwargs, tensorboard_log, verbose,
                                                device, support_multi_env, create_eval_env, monitor_wrapper,
                                                seed, use_sde, sde_sample_freq)
         self.buffer_size = buffer_size
         self.batch_size = batch_size
         self.learning_starts = learning_starts
-        self.actor = None
+        self.actor = None  # type: Optional[th.nn.Module]
         self.replay_buffer = None  # type: Optional[ReplayBuffer]
         # Update policy keyword arguments
         if sde_support:
@@ -752,7 +766,7 @@ class OffPolicyRLModel(BaseRLModel):
             self.replay_buffer = pickle.load(file_handler)
         assert isinstance(self.replay_buffer, ReplayBuffer), 'The replay buffer must inherit from ReplayBuffer class'
 
-    def collect_rollouts(self,
+    def collect_rollouts(self,  # noqa: C901
                          env: VecEnv,
                          # Type hint as string to avoid circular import
                          callback: 'BaseCallback',
@@ -873,22 +887,23 @@ class OffPolicyRLModel(BaseRLModel):
                 if action_noise is not None:
                     action_noise.reset()
 
-                # Display training infos
+                # Log training infos
-                if self.verbose >= 1 and log_interval is not None and self._episode_num % log_interval == 0:
+                if log_interval is not None and self._episode_num % log_interval == 0:
                     fps = int(self.num_timesteps / (time.time() - self.start_time))
-                    logger.logkv("episodes", self._episode_num)
+                    logger.record("time/episodes", self._episode_num, exclude="tensorboard")
                     if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
-                        logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
+                        logger.record('rollout/ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
-                        logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
+                        logger.record('rollout/ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
-                    logger.logkv("fps", fps)
+                    logger.record("time/fps", fps)
-                    logger.logkv('time_elapsed', int(time.time() - self.start_time))
+                    logger.record('time/time_elapsed', int(time.time() - self.start_time), exclude="tensorboard")
-                    logger.logkv("total timesteps", self.num_timesteps)
+                    logger.record("time/total timesteps", self.num_timesteps, exclude="tensorboard")
                     if self.use_sde:
-                        logger.logkv("std", (self.actor.get_std()).mean().item())
+                        logger.record("train/std", (self.actor.get_std()).mean().item())
 
                     if len(self.ep_success_buffer) > 0:
-                        logger.logkv('success rate', self.safe_mean(self.ep_success_buffer))
+                        logger.record('rollout/success rate', self.safe_mean(self.ep_success_buffer))
-                    logger.dumpkvs()
+                    # Pass the number of timesteps for tensorboard
+                    logger.dump(step=self.num_timesteps)
 
         mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
 

stable_baselines3/common/callbacks.py

@@ -9,7 +9,7 @@ import numpy as np
 
 from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv, sync_envs_normalization
 from stable_baselines3.common.evaluation import evaluate_policy
-from stable_baselines3.common.logger import Logger
+from stable_baselines3.common import logger
 
 if typing.TYPE_CHECKING:
     from stable_baselines3.common.base_class import BaseRLModel  # pytype: disable=pyi-error
@@ -34,7 +34,7 @@ class BaseCallback(ABC):
         self.verbose = verbose
         self.locals = None  # type: Optional[Dict[str, Any]]
         self.globals = None  # type: Optional[Dict[str, Any]]
-        self.logger = None  # type: Optional[Logger]
+        self.logger = None
         # Sometimes, for event callback, it is useful
         # to have access to the parent object
         self.parent = None  # type: Optional[BaseCallback]
@@ -47,7 +47,7 @@ class BaseCallback(ABC):
         """
         self.model = model
         self.training_env = model.get_env()
-        self.logger = Logger.CURRENT
+        self.logger = logger
         self._init_callback()
 
     def _init_callback(self) -> None:
@@ -313,6 +313,9 @@ class EvalCallback(EventCallback):
                 print(f"Eval num_timesteps={self.num_timesteps}, "
                       f"episode_reward={mean_reward:.2f} +/- {std_reward:.2f}")
                 print(f"Episode length: {mean_ep_length:.2f} +/- {std_ep_length:.2f}")
+            # Add to current Logger
+            self.logger.record('eval/mean_reward', float(mean_reward))
+            self.logger.record('eval/mean_ep_length', mean_ep_length)
 
             if mean_reward > self.best_mean_reward:
                 if self.verbose > 0:

stable_baselines3/common/logger.py

@@ -5,9 +5,15 @@ import os
 import tempfile
 import warnings
 from collections import defaultdict
-from typing import Dict, List, TextIO, Union, Any, Optional
+from typing import Dict, List, TextIO, Union, Any, Optional, Tuple
 
 import pandas
+import numpy as np
+import torch as th
+try:
+    from torch.utils.tensorboard import SummaryWriter
+except ImportError:
+    SummaryWriter = None
 
 DEBUG = 10
 INFO = 20
@@ -21,11 +27,14 @@ class KVWriter(object):
     Key Value writer
     """
 
-    def writekvs(self, kvs: Dict) -> None:
+    def write(self, key_values: Dict[str, Any],
+              key_excluded: Dict[str, Union[str, Tuple[str, ...]]], step: int = 0) -> None:
         """
-        write a dictionary to file
+        Write a dictionary to file
 
-        :param kvs: (dict)
+        :param key_values: (dict)
+        :param key_excluded: (dict)
+        :param step: (int)
         """
         raise NotImplementedError
 
@@ -41,11 +50,11 @@ class SeqWriter(object):
     sequence writer
     """
 
-    def writeseq(self, seq: List):
+    def write_sequence(self, sequence: List):
         """
-        write an array to file
+        write_sequence an array to file
 
-        :param seq: (list)
+        :param sequence: (list)
         """
         raise NotImplementedError
 
@@ -65,16 +74,29 @@ class HumanOutputFormat(KVWriter, SeqWriter):
             self.file = filename_or_file
             self.own_file = False
 
-    def writekvs(self, kvs: Dict) -> None:
+    def write(self, key_values: Dict, key_excluded: Dict, step: int = 0) -> None:
         # Create strings for printing
         key2str = {}
-        for (key, val) in sorted(kvs.items()):
+        tag = None
-            if isinstance(val, float):
+        for (key, value), (_, excluded) in zip(sorted(key_values.items()), sorted(key_excluded.items())):
+
+            if excluded is not None and 'stdout' in excluded:
+                continue
+
+            if isinstance(value, float):
                 # Align left
-                val_str = f'{val:<8.3g}'
+                value_str = f'{value:<8.3g}'
             else:
-                val_str = str(val)
+                value_str = str(value)
-            key2str[self._truncate(key)] = self._truncate(val_str)
+
+            if key.find('/') > 0:  # Find tag and add it to the dict
+                tag = key[:key.find('/') + 1]
+                key2str[self._truncate(tag)] = ''
+            # Remove tag from key
+            if tag is not None and tag in key:
+                key = str('   ' + key[len(tag):])
+
+            key2str[self._truncate(key)] = self._truncate(value_str)
 
         # Find max widths
         if len(key2str) == 0:
@@ -87,10 +109,10 @@ class HumanOutputFormat(KVWriter, SeqWriter):
         # Write out the data
         dashes = '-' * (key_width + val_width + 7)
         lines = [dashes]
-        for (key, val) in sorted(key2str.items()):
+        for key, value in key2str.items():
             key_space = ' ' * (key_width - len(key))
-            val_space = ' ' * (val_width - len(val))
+            val_space = ' ' * (val_width - len(value))
-            lines.append(f"| {key}{key_space} | {val}{val_space} |")
+            lines.append(f"| {key}{key_space} | {value}{val_space} |")
         lines.append(dashes)
         self.file.write('\n'.join(lines) + '\n')
 
@@ -98,14 +120,14 @@ class HumanOutputFormat(KVWriter, SeqWriter):
         self.file.flush()
 
     @classmethod
-    def _truncate(cls, string: str) -> str:
+    def _truncate(cls, string: str, max_length: int = 23) -> str:
-        return string[:20] + '...' if len(string) > 23 else string
+        return string[:max_length - 3] + '...' if len(string) > max_length else string
 
-    def writeseq(self, seq: List) -> None:
+    def write_sequence(self, sequence: List) -> None:
-        seq = list(seq)
+        sequence = list(sequence)
-        for (i, elem) in enumerate(seq):
+        for i, elem in enumerate(sequence):
             self.file.write(elem)
-            if i < len(seq) - 1:  # add space unless this is the last one
+            if i < len(sequence) - 1:  # add space unless this is the last one
                 self.file.write(' ')
         self.file.write('\n')
         self.file.flush()
@@ -127,22 +149,28 @@ class JSONOutputFormat(KVWriter):
         """
         self.file = open(filename, 'wt')
 
-    def writekvs(self, kvs: Dict) -> None:
+    def write(self, key_values: Dict[str, Any],
-        for key, value in sorted(kvs.items()):
+              key_excluded: Dict[str, Union[str, Tuple[str, ...]]], step: int = 0) -> None:
+        for (key, value), (_, excluded) in zip(sorted(key_values.items()), sorted(key_excluded.items())):
+
+            if excluded is not None and 'json' in excluded:
+                continue
+
             if hasattr(value, 'dtype'):
                 if value.shape == () or len(value) == 1:
                     # if value is a dimensionless numpy array or of length 1, serialize as a float
-                    kvs[key] = float(value)
+                    key_values[key] = float(value)
                 else:
                     # otherwise, a value is a numpy array, serialize as a list or nested lists
-                    kvs[key] = value.tolist()
+                    key_values[key] = value.tolist()
-        self.file.write(json.dumps(kvs) + '\n')
+        self.file.write(json.dumps(key_values) + '\n')
         self.file.flush()
 
     def close(self) -> None:
         """
         closes the file
         """
+
         self.file.close()
 
 
@@ -153,13 +181,15 @@ class CSVOutputFormat(KVWriter):
 
         :param filename: (str) the file to write the log to
         """
+
         self.file = open(filename, 'w+t')
         self.keys = []
-        self.sep = ','
+        self.separator = ','
 
-    def writekvs(self, kvs: Dict) -> None:
+    def write(self, key_values: Dict[str, Any],
+              key_excluded: Dict[str, Union[str, Tuple[str, ...]]], step: int = 0) -> None:
         # Add our current row to the history
-        extra_keys = kvs.keys() - self.keys
+        extra_keys = key_values.keys() - self.keys
         if extra_keys:
             self.keys.extend(extra_keys)
             self.file.seek(0)
@@ -172,12 +202,12 @@ class CSVOutputFormat(KVWriter):
             self.file.write('\n')
             for line in lines[1:]:
                 self.file.write(line[:-1])
-                self.file.write(self.sep * len(extra_keys))
+                self.file.write(self.separator * len(extra_keys))
                 self.file.write('\n')
         for i, key in enumerate(self.keys):
             if i > 0:
                 self.file.write(',')
-            value = kvs.get(key)
+            value = key_values.get(key)
             if value is not None:
                 self.file.write(str(value))
         self.file.write('\n')
@@ -190,25 +220,49 @@ class CSVOutputFormat(KVWriter):
         self.file.close()
 
 
-def valid_float_value(value: Any) -> bool:
+class TensorBoardOutputFormat(KVWriter):
-    """
+    def __init__(self, folder: str):
-    Returns True if the value can be successfully cast into a float
+        """
+        Dumps key/value pairs into TensorBoard's numeric format.
 
-    :param value: (Any) the value to check
+        :param folder: (str) the folder to write the log to
-    :return: (bool)
+        """
-    """
+        assert SummaryWriter is not None, ("tensorboard is not installed, you can use "
-    try:
+                                           "pip install tensorboard to do so")
-        float(value)
+        self.writer = SummaryWriter(log_dir=folder)
-        return True
+
-    except TypeError:
+    def write(self, key_values: Dict[str, Any],
-        return False
+              key_excluded: Dict[str, Union[str, Tuple[str, ...]]], step: int = 0) -> None:
+
+        for (key, value), (_, excluded) in zip(sorted(key_values.items()),
+                                               sorted(key_excluded.items())):
+
+            if excluded is not None and 'tensorboard' in excluded:
+                continue
+
+            if isinstance(value, np.ScalarType):
+                self.writer.add_scalar(key, value, step)
+
+            if isinstance(value, th.Tensor):
+                self.writer.add_histogram(key, value, step)
+
+        # Flush the output to the file
+        self.writer.flush()
+
+    def close(self) -> None:
+        """
+        closes the file
+        """
+        if self.writer:
+            self.writer.close()
+            self.writer = None
 
 
 def make_output_format(_format: str, log_dir: str, log_suffix: str = '') -> KVWriter:
     """
     return a logger for the requested format
 
-    :param _format: (str) the requested format to log to ('stdout', 'log', 'json' or 'csv')
+    :param _format: (str) the requested format to log to ('stdout', 'log', 'json' or 'csv' or 'tensorboard')
     :param log_dir: (str) the logging directory
     :param log_suffix: (str) the suffix for the log file
     :return: (KVWriter) the logger
@@ -222,6 +276,8 @@ def make_output_format(_format: str, log_dir: str, log_suffix: str = '') -> KVWr
         return JSONOutputFormat(os.path.join(log_dir, f'progress{log_suffix}.json'))
     elif _format == 'csv':
         return CSVOutputFormat(os.path.join(log_dir, f'progress{log_suffix}.csv'))
+    elif _format == 'tensorboard':
+        return TensorBoardOutputFormat(log_dir)
     else:
         raise ValueError(f'Unknown format specified: {_format}')
 
@@ -230,52 +286,56 @@ def make_output_format(_format: str, log_dir: str, log_suffix: str = '') -> KVWr
 # API
 # ================================================================
 
-def logkv(key: Any, val: Any) -> None:
+def record(key: str, value: Any,
+           exclude: Optional[Union[str, Tuple[str, ...]]] = None) -> None:
     """
     Log a value of some diagnostic
     Call this once for each diagnostic quantity, each iteration
     If called many times, last value will be used.
 
     :param key: (Any) save to log this key
-    :param val: (Any) save to log this value
+    :param value: (Any) save to log this value
+    :param exclude: (str or tuple) outputs to be excluded
     """
-    Logger.CURRENT.logkv(key, val)
+    Logger.CURRENT.record(key, value, exclude)
 
 
-def logkv_mean(key: Any, val: Union[int, float]) -> None:
+def record_mean(key: str, value: Union[int, float],
+                exclude: Optional[Union[str, Tuple[str, ...]]] = None) -> None:
     """
-    The same as logkv(), but if called many times, values averaged.
+    The same as record(), but if called many times, values averaged.
 
     :param key: (Any) save to log this key
-    :param val: (Number) save to log this value
+    :param value: (Number) save to log this value
+    :param exclude: (str or tuple) outputs to be excluded
     """
-    Logger.CURRENT.logkv_mean(key, val)
+    Logger.CURRENT.record_mean(key, value, exclude)
 
 
-def logkvs(key_values: Dict) -> None:
+def record_dict(key_values: Dict[str, Any]) -> None:
     """
-    Log a dictionary of key-value pairs
+    Log a dictionary of key-value pairs.
 
     :param key_values: (dict) the list of keys and values to save to log
     """
     for key, value in key_values.items():
-        logkv(key, value)
+        record(key, value)
 
 
-def dumpkvs() -> None:
+def dump(step: int = 0) -> None:
     """
     Write all of the diagnostics from the current iteration
     """
-    Logger.CURRENT.dumpkvs()
+    Logger.CURRENT.dump(step)
 
 
-def getkvs() -> Dict:
+def get_log_dict() -> Dict:
     """
     get the key values logs
 
     :return: (dict) the logged values
     """
-    return Logger.CURRENT.name2val
+    return Logger.CURRENT.name_to_value
 
 
 def log(*args, level: int = INFO) -> None:
@@ -363,8 +423,8 @@ def get_dir() -> str:
     return Logger.CURRENT.get_dir()
 
 
-record_tabular = logkv
+record_tabular = record
-dump_tabular = dumpkvs
+dump_tabular = dump
 
 
 # ================================================================
@@ -384,50 +444,59 @@ class Logger(object):
         :param folder: (str) the logging location
         :param output_formats: ([str]) the list of output format
         """
-        self.name2val = defaultdict(float)  # values this iteration
+        self.name_to_value = defaultdict(float)  # values this iteration
-        self.name2cnt = defaultdict(int)
+        self.name_to_count = defaultdict(int)
+        self.name_to_excluded = defaultdict(str)
         self.level = INFO
         self.dir = folder
         self.output_formats = output_formats
 
     # Logging API, forwarded
     # ----------------------------------------
-    def logkv(self, key: Any, val: Any) -> None:
+    def record(self, key: str, value: Any,
+               exclude: Optional[Union[str, Tuple[str, ...]]] = None) -> None:
         """
         Log a value of some diagnostic
         Call this once for each diagnostic quantity, each iteration
         If called many times, last value will be used.
 
         :param key: (Any) save to log this key
-        :param val: (Any) save to log this value
+        :param value: (Any) save to log this value
+        :param exclude: (str or tuple) outputs to be excluded
         """
-        self.name2val[key] = val
+        self.name_to_value[key] = value
+        self.name_to_excluded[key] = exclude
 
-    def logkv_mean(self, key: Any, val: Any) -> None:
+    def record_mean(self, key: str, value: Any,
+                    exclude: Optional[Union[str, Tuple[str, ...]]] = None) -> None:
         """
-        The same as logkv(), but if called many times, values averaged.
+        The same as record(), but if called many times, values averaged.
 
         :param key: (Any) save to log this key
-        :param val: (Number) save to log this value
+        :param value: (Number) save to log this value
+        :param exclude: (str or tuple) outputs to be excluded
         """
-        if val is None:
+        if value is None:
-            self.name2val[key] = None
+            self.name_to_value[key] = None
             return
-        oldval, cnt = self.name2val[key], self.name2cnt[key]
+        old_val, count = self.name_to_value[key], self.name_to_count[key]
-        self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1)
+        self.name_to_value[key] = old_val * count / (count + 1) + value / (count + 1)
-        self.name2cnt[key] = cnt + 1
+        self.name_to_count[key] = count + 1
+        self.name_to_excluded[key] = exclude
 
-    def dumpkvs(self) -> None:
+    def dump(self, step: int = 0) -> None:
         """
         Write all of the diagnostics from the current iteration
         """
         if self.level == DISABLED:
             return
-        for fmt in self.output_formats:
+        for _format in self.output_formats:
-            if isinstance(fmt, KVWriter):
+            if isinstance(_format, KVWriter):
-                fmt.writekvs(self.name2val)
+                _format.write(self.name_to_value, self.name_to_excluded, step)
-        self.name2val.clear()
+
-        self.name2cnt.clear()
+        self.name_to_value.clear()
+        self.name_to_count.clear()
+        self.name_to_excluded.clear()
 
     def log(self, *args, level: int = INFO) -> None:
         """
@@ -466,8 +535,8 @@ class Logger(object):
         """
         closes the file
         """
-        for fmt in self.output_formats:
+        for _format in self.output_formats:
-            fmt.close()
+            _format.close()
 
     # Misc
     # ----------------------------------------
@@ -477,37 +546,37 @@ class Logger(object):
 
         :param args: (list) the arguments to log
         """
-        for fmt in self.output_formats:
+        for _format in self.output_formats:
-            if isinstance(fmt, SeqWriter):
+            if isinstance(_format, SeqWriter):
-                fmt.writeseq(map(str, args))
+                _format.write_sequence(map(str, args))
 
 
 # Initialize logger
 Logger.DEFAULT = Logger.CURRENT = Logger(folder=None, output_formats=[HumanOutputFormat(sys.stdout)])
 
 
-def configure(folder: Optional[str] = None, format_strs: Optional[List[str]] = None) -> None:
+def configure(folder: Optional[str] = None, format_strings: Optional[List[str]] = None) -> None:
     """
     configure the current logger
 
     :param folder: (Optional[str]) the save location
-        (if None, $BASELINES_LOGDIR, if still None, tempdir/baselines-[date & time])
+        (if None, $SB3_LOGDIR, if still None, tempdir/baselines-[date & time])
-    :param format_strs: (Optional[List[str]]) the output logging format
+    :param format_strings: (Optional[List[str]]) the output logging format
-        (if None, $BASELINES_LOG_FORMAT, if still None, ['stdout', 'log', 'csv'])
+        (if None, $SB3_LOG_FORMAT, if still None, ['stdout', 'log', 'csv'])
     """
     if folder is None:
-        folder = os.getenv('BASELINES_LOGDIR')
+        folder = os.getenv('SB3_LOGDIR')
     if folder is None:
-        folder = os.path.join(tempfile.gettempdir(), datetime.datetime.now().strftime("baselines-%Y-%m-%d-%H-%M-%S-%f"))
+        folder = os.path.join(tempfile.gettempdir(), datetime.datetime.now().strftime("SB3-%Y-%m-%d-%H-%M-%S-%f"))
     assert isinstance(folder, str)
     os.makedirs(folder, exist_ok=True)
 
     log_suffix = ''
-    if format_strs is None:
+    if format_strings is None:
-        format_strs = os.getenv('BASELINES_LOG_FORMAT', 'stdout,log,csv').split(',')
+        format_strings = os.getenv('SB3_LOG_FORMAT', 'stdout,log,csv').split(',')
 
-    format_strs = filter(None, format_strs)
+    format_strings = filter(None, format_strings)
-    output_formats = [make_output_format(f, folder, log_suffix) for f in format_strs]
+    output_formats = [make_output_format(f, folder, log_suffix) for f in format_strings]
 
     Logger.CURRENT = Logger(folder=folder, output_formats=output_formats)
     log(f'Logging to {folder}')
@@ -524,28 +593,28 @@ def reset() -> None:
 
 
 class ScopedConfigure(object):
-    def __init__(self, folder: Optional[str] = None, format_strs: Optional[List[str]] = None):
+    def __init__(self, folder: Optional[str] = None, format_strings: Optional[List[str]] = None):
         """
         Class for using context manager while logging
 
         usage:
-        with ScopedConfigure(folder=None, format_strs=None):
+        with ScopedConfigure(folder=None, format_strings=None):
             {code}
 
         :param folder: (str) the logging folder
-        :param format_strs: ([str]) the list of output logging format
+        :param format_strings: ([str]) the list of output logging format
         """
         self.dir = folder
-        self.format_strs = format_strs
+        self.format_strings = format_strings
-        self.prevlogger = None
+        self.prev_logger = None
 
     def __enter__(self) -> None:
-        self.prevlogger = Logger.CURRENT
+        self.prev_logger = Logger.CURRENT
-        configure(folder=self.dir, format_strs=self.format_strs)
+        configure(folder=self.dir, format_strings=self.format_strings)
 
     def __exit__(self, *args) -> None:
         Logger.CURRENT.close()
-        Logger.CURRENT = self.prevlogger
+        Logger.CURRENT = self.prev_logger
 
 
 # ================================================================

stable_baselines3/common/utils.py

@@ -1,8 +1,17 @@
-from typing import Callable, Union
+from typing import Callable, Union, Optional
 import random
 
+import os
+import glob
 import numpy as np
 import torch as th
+# Check if tensorboard is available for pytorch
+try:
+    from torch.utils.tensorboard import SummaryWriter
+except ImportError:
+    SummaryWriter = None
+
+from stable_baselines3.common import logger
 
 
 def set_random_seed(seed: int, using_cuda: bool = False) -> None:
@@ -110,3 +119,42 @@ def get_device(device: Union[th.device, str] = 'auto') -> th.device:
         return th.device('cpu')
 
     return device
+
+
+def get_latest_run_id(log_path: Optional[str] = None, log_name: str = '') -> int:
+    """
+    Returns the latest run number for the given log name and log path,
+    by finding the greatest number in the directories.
+
+    :return: (int) latest run number
+    """
+    max_run_id = 0
+    for path in glob.glob(f"{log_path}/{log_name}_[0-9]*"):
+        file_name = path.split(os.sep)[-1]
+        ext = file_name.split("_")[-1]
+        if log_name == "_".join(file_name.split("_")[:-1]) and ext.isdigit() and int(ext) > max_run_id:
+            max_run_id = int(ext)
+    return max_run_id
+
+
+def configure_logger(verbose: int = 0, tensorboard_log: Optional[str] = None,
+                     tb_log_name: str = '', reset_num_timesteps: bool = True) -> None:
+    """
+    Configure the logger's outputs.
+
+    :param verbose: (int) the verbosity level: 0 no output, 1 info, 2 debug
+    :param tensorboard_log: (str) the log location for tensorboard (if None, no logging)
+    :param tb_log_name: (str) tensorboard log
+    """
+    if tensorboard_log is not None and SummaryWriter is not None:
+        latest_run_id = get_latest_run_id(tensorboard_log, tb_log_name)
+        if not reset_num_timesteps:
+            # Continue training in the same directory
+            latest_run_id -= 1
+        save_path = os.path.join(tensorboard_log, f"{tb_log_name}_{latest_run_id + 1}")
+        if verbose >= 1:
+            logger.configure(save_path, ["stdout", "tensorboard"])
+        else:
+            logger.configure(save_path, ["tensorboard"])
+    elif verbose == 0:
+        logger.configure(format_strings=[""])

stable_baselines3/ppo/ppo.py

@@ -5,13 +5,6 @@ import gym
 from gym import spaces
 import torch as th
 import torch.nn.functional as F
-
-# Check if tensorboard is available for pytorch
-# TODO: finish tensorboard integration
-# try:
-#     from torch.utils.tensorboard import SummaryWriter
-# except ImportError:
-#     SummaryWriter = None
 import numpy as np
 
 from stable_baselines3.common import logger
@@ -95,10 +88,11 @@ class PPO(BaseRLModel):
                  policy_kwargs: Optional[Dict[str, Any]] = None,
                  verbose: int = 0,
                  seed: Optional[int] = None,
-                 device: Union[th.device, str] = 'auto',
+                 device: Union[th.device, str] = "auto",
                  _init_setup_model: bool = True):
 
-        super(PPO, self).__init__(policy, env, PPOPolicy, learning_rate, policy_kwargs=policy_kwargs,
+        super(PPO, self).__init__(policy, env, PPOPolicy, learning_rate,
+                                  policy_kwargs=policy_kwargs, tensorboard_log=tensorboard_log,
                                   verbose=verbose, device=device, use_sde=use_sde, sde_sample_freq=sde_sample_freq,
                                   create_eval_env=create_eval_env, support_multi_env=True, seed=seed)
 
@@ -114,7 +108,6 @@ class PPO(BaseRLModel):
         self.max_grad_norm = max_grad_norm
         self.rollout_buffer = None
         self.target_kl = target_kl
-        self.tensorboard_log = tensorboard_log
         self.tb_writer = None
 
         if _init_setup_model:
@@ -136,8 +129,8 @@ class PPO(BaseRLModel):
         self.clip_range = get_schedule_fn(self.clip_range)
         if self.clip_range_vf is not None:
             if isinstance(self.clip_range_vf, (float, int)):
-                assert self.clip_range_vf > 0, ('`clip_range_vf` must be positive, '
+                assert self.clip_range_vf > 0, ("`clip_range_vf` must be positive, "
-                                                'pass `None` to deactivate vf clipping')
+                                                "pass `None` to deactivate vf clipping")
 
             self.clip_range_vf = get_schedule_fn(self.clip_range_vf)
 
@@ -231,6 +224,7 @@ class PPO(BaseRLModel):
 
                 # ratio between old and new policy, should be one at the first iteration
                 ratio = th.exp(log_prob - rollout_data.old_log_prob)
+
                 # clipped surrogate loss
                 policy_loss_1 = advantages * ratio
                 policy_loss_2 = advantages * th.clamp(ratio, 1 - clip_range, 1 + clip_range)
@@ -282,19 +276,21 @@ class PPO(BaseRLModel):
         explained_var = explained_variance(self.rollout_buffer.returns.flatten(),
                                            self.rollout_buffer.values.flatten())
 
-        logger.logkv("n_updates", self._n_updates)
+        # Logs
-        logger.logkv("clip_fraction", np.mean(clip_fraction))
+        logger.record("train/entropy_loss", np.mean(entropy_losses))
-        logger.logkv("clip_range", clip_range)
+        logger.record("train/policy_gradient_loss", np.mean(pg_losses))
-        if self.clip_range_vf is not None:
+        logger.record("train/value_loss", np.mean(value_losses))
-            logger.logkv("clip_range_vf", clip_range_vf)
+        logger.record("train/approx_kl", np.mean(approx_kl_divs))
+        logger.record("train/clip_fraction", np.mean(clip_fraction))
+        logger.record("train/loss", loss.item())
+        logger.record("train/explained_variance", explained_var)
+        if hasattr(self.policy, "log_std"):
+            logger.record("train/std", th.exp(self.policy.log_std).mean().item())
 
-        logger.logkv("approx_kl", np.mean(approx_kl_divs))
+        logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
-        logger.logkv("explained_variance", explained_var)
+        logger.record("train/clip_range", clip_range)
-        logger.logkv("entropy_loss", np.mean(entropy_losses))
+        if self.clip_range_vf is not None:
-        logger.logkv("policy_gradient_loss", np.mean(pg_losses))
+            logger.record("train/clip_range_vf", clip_range_vf)
-        logger.logkv("value_loss", np.mean(value_losses))
-        if hasattr(self.policy, 'log_std'):
-            logger.logkv("std", th.exp(self.policy.log_std).mean().item())
 
     def learn(self,
               total_timesteps: int,
@@ -305,15 +301,12 @@ class PPO(BaseRLModel):
               n_eval_episodes: int = 5,
               tb_log_name: str = "PPO",
               eval_log_path: Optional[str] = None,
-              reset_num_timesteps: bool = True) -> 'PPO':
+              reset_num_timesteps: bool = True) -> "PPO":
 
         iteration = 0
-        callback = self._setup_learn(eval_env, callback, eval_freq,
+        total_timesteps, callback = self._setup_learn(total_timesteps, eval_env, callback, eval_freq,
-                                     n_eval_episodes, eval_log_path, reset_num_timesteps)
+                                                      n_eval_episodes, eval_log_path, reset_num_timesteps,
-
+                                                      tb_log_name)
-        # if self.tensorboard_log is not None and SummaryWriter is not None:
-        #     self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name))
-
         callback.on_training_start(locals(), globals())
 
         while self.num_timesteps < total_timesteps:
@@ -328,24 +321,22 @@ class PPO(BaseRLModel):
             iteration += 1
             self._update_current_progress(self.num_timesteps, total_timesteps)
 
-            # Display training infos
+            # Log training infos
-            if self.verbose >= 1 and log_interval is not None and iteration % log_interval == 0:
+            if log_interval is not None and iteration % log_interval == 0:
                 fps = int(self.num_timesteps / (time.time() - self.start_time))
-                logger.logkv("iterations", iteration)
+                logger.record("time/iterations", iteration, exclude="tensorboard")
                 if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
-                    logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
+                    logger.record("rollout/ep_rew_mean",
-                    logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
+                                  self.safe_mean([ep_info["r"] for ep_info in self.ep_info_buffer]))
-                logger.logkv("fps", fps)
+                    logger.record("rollout/ep_len_mean",
-                logger.logkv('time_elapsed', int(time.time() - self.start_time))
+                                  self.safe_mean([ep_info["l"] for ep_info in self.ep_info_buffer]))
-                logger.logkv("total timesteps", self.num_timesteps)
+                logger.record("time/fps", fps)
-                logger.dumpkvs()
+                logger.record("time/time_elapsed", int(time.time() - self.start_time), exclude="tensorboard")
+                logger.record("time/total timesteps", self.num_timesteps, exclude="tensorboard")
+                logger.dump(step=self.num_timesteps)
 
             self.train(self.n_epochs, batch_size=self.batch_size)
 
-            # For tensorboard integration
-            # if self.tb_writer is not None:
-            #     self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
-
         callback.on_training_end()
 
         return self

stable_baselines3/sac/sac.py

@@ -1,5 +1,4 @@
 from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
-
 import torch as th
 import torch.nn.functional as F
 import numpy as np
@@ -90,7 +89,7 @@ class SAC(OffPolicyRLModel):
 
         super(SAC, self).__init__(policy, env, SACPolicy, learning_rate,
                                   buffer_size, learning_starts, batch_size,
-                                  policy_kwargs, verbose, device,
+                                  policy_kwargs, tensorboard_log, verbose, device,
                                   create_eval_env=create_eval_env, seed=seed,
                                   use_sde=use_sde, sde_sample_freq=sde_sample_freq,
                                   use_sde_at_warmup=use_sde_at_warmup)
@@ -237,12 +236,12 @@ class SAC(OffPolicyRLModel):
 
         self._n_updates += gradient_steps
 
-        logger.logkv("n_updates", self._n_updates)
+        logger.record("train/n_updates", self._n_updates, exclude='tensorboard')
-        logger.logkv("ent_coef", np.mean(ent_coefs))
+        logger.record("train/ent_coef", np.mean(ent_coefs))
-        logger.logkv("actor_loss", np.mean(actor_losses))
+        logger.record("train/actor_loss", np.mean(actor_losses))
-        logger.logkv("critic_loss", np.mean(critic_losses))
+        logger.record("train/critic_loss", np.mean(critic_losses))
         if len(ent_coef_losses) > 0:
-            logger.logkv("ent_coef_loss", np.mean(ent_coef_losses))
+            logger.record("train/ent_coef_loss", np.mean(ent_coef_losses))
 
     def learn(self,
               total_timesteps: int,
@@ -255,8 +254,9 @@ class SAC(OffPolicyRLModel):
               eval_log_path: Optional[str] = None,
               reset_num_timesteps: bool = True) -> OffPolicyRLModel:
 
-        callback = self._setup_learn(eval_env, callback, eval_freq,
+        total_timesteps, callback = self._setup_learn(total_timesteps, eval_env, callback, eval_freq,
-                                     n_eval_episodes, eval_log_path, reset_num_timesteps)
+                                                      n_eval_episodes, eval_log_path, reset_num_timesteps,
+                                                      tb_log_name)
         callback.on_training_start(locals(), globals())
 
         while self.num_timesteps < total_timesteps:

stable_baselines3/td3/td3.py

@@ -74,7 +74,7 @@ class TD3(OffPolicyRLModel):
 
         super(TD3, self).__init__(policy, env, TD3Policy, learning_rate,
                                   buffer_size, learning_starts, batch_size,
-                                  policy_kwargs, verbose, device,
+                                  policy_kwargs, tensorboard_log, verbose, device,
                                   create_eval_env=create_eval_env, seed=seed,
                                   sde_support=False)
 
@@ -152,7 +152,7 @@ class TD3(OffPolicyRLModel):
                     target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
 
         self._n_updates += gradient_steps
-        logger.logkv("n_updates", self._n_updates)
+        logger.record("train/n_updates", self._n_updates, exclude='tensorboard')
 
     def learn(self,
               total_timesteps: int,
@@ -165,9 +165,9 @@ class TD3(OffPolicyRLModel):
               eval_log_path: Optional[str] = None,
               reset_num_timesteps: bool = True) -> OffPolicyRLModel:
 
-        callback = self._setup_learn(eval_env, callback, eval_freq,
+        total_timesteps, callback = self._setup_learn(total_timesteps, eval_env, callback, eval_freq,
-                                     n_eval_episodes, eval_log_path, reset_num_timesteps)
+                                                      n_eval_episodes, eval_log_path, reset_num_timesteps,
-
+                                                      tb_log_name)
         callback.on_training_start(locals(), globals())
 
         while self.num_timesteps < total_timesteps:

stable_baselines3/version.txt

@@ -1 +1 @@
-0.6.0a10
+0.6.0a11

tests/test_logger.py

@@ -1,12 +1,9 @@
-import os
-import shutil
-
 import pytest
 import numpy as np
 
 from stable_baselines3.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure,
-                                             info, debug, set_level, configure, logkv, logkvs,
+                                             info, debug, set_level, configure, record, record_dict,
-                                             dumpkvs, logkv_mean, warn, error, reset)
+                                             dump, record_mean, warn, error, reset)
 
 KEY_VALUES = {
     "test": 1,
@@ -18,10 +15,12 @@ KEY_VALUES = {
     "g": np.array([[[1]]]),
 }
 
-LOG_DIR = '/tmp/stable_baselines3/'
+KEY_EXCLUDED = {}
+for key in KEY_VALUES.keys():
+    KEY_EXCLUDED[key] = None
 
 
-def test_main():
+def test_main(tmp_path):
     """
     tests for the logger module
     """
@@ -29,55 +28,56 @@ def test_main():
     debug("shouldn't appear")
     set_level(DEBUG)
     debug("should appear")
-    folder = "/tmp/testlogging"
+    configure(folder=str(tmp_path))
-    if os.path.exists(folder):
+    record("a", 3)
-        shutil.rmtree(folder)
+    record("b", 2.5)
-    configure(folder=folder)
+    dump()
-    logkv("a", 3)
+    record("b", -2.5)
-    logkv("b", 2.5)
+    record("a", 5.5)
-    dumpkvs()
+    dump()
-    logkv("b", -2.5)
-    logkv("a", 5.5)
-    dumpkvs()
     info("^^^ should see a = 5.5")
-    logkv_mean("b", -22.5)
+    record_mean("b", -22.5)
-    logkv_mean("b", -44.4)
+    record_mean("b", -44.4)
-    logkv("a", 5.5)
+    record("a", 5.5)
-    dumpkvs()
+    dump()
     with ScopedConfigure(None, None):
         info("^^^ should see b = 33.3")
 
-    with ScopedConfigure("/tmp/test-logger/", ["json"]):
+    with ScopedConfigure(str(tmp_path / "test-logger"), ["json"]):
-        logkv("b", -2.5)
+        record("b", -2.5)
-        dumpkvs()
+        dump()
 
     reset()
-    logkv("a", "longasslongasslongasslongasslongasslongassvalue")
+    record("a", "longasslongasslongasslongasslongasslongassvalue")
-    dumpkvs()
+    dump()
     warn("hey")
     error("oh")
-    logkvs({"test": 1})
+    record_dict({"test": 1})
 
 
-@pytest.mark.parametrize('_format', ['stdout', 'log', 'json', 'csv'])
+@pytest.mark.parametrize('_format', ['stdout', 'log', 'json', 'csv', 'tensorboard'])
-def test_make_output(_format):
+def test_make_output(tmp_path, _format):
     """
     test make output
 
     :param _format: (str) output format
     """
-    writer = make_output_format(_format, LOG_DIR)
+    if _format == 'tensorboard':
-    writer.writekvs(KEY_VALUES)
+        # Skip if no tensorboard installed
+        pytest.importorskip("tensorboard")
+
+    writer = make_output_format(_format, tmp_path)
+    writer.write(KEY_VALUES, KEY_EXCLUDED)
     if _format == "csv":
-        read_csv(LOG_DIR + 'progress.csv')
+        read_csv(tmp_path / 'progress.csv')
     elif _format == 'json':
-        read_json(LOG_DIR + 'progress.json')
+        read_json(tmp_path / 'progress.json')
     writer.close()
 
 
-def test_make_output_fail():
+def test_make_output_fail(tmp_path):
     """
     test value error on logger
     """
     with pytest.raises(ValueError):
-        make_output_format('dummy_format', LOG_DIR)
+        make_output_format('dummy_format', tmp_path)

tests/test_tensorboard.py

@@ -0,0 +1,36 @@
+import os
+import pytest
+
+from stable_baselines3 import A2C, PPO, SAC, TD3
+
+MODEL_DICT = {
+    'a2c': (A2C, 'CartPole-v1'),
+    'ppo': (PPO, 'CartPole-v1'),
+    'sac': (SAC, 'Pendulum-v0'),
+    'td3': (TD3, 'Pendulum-v0'),
+}
+
+N_STEPS = 100
+
+
+@pytest.mark.parametrize("model_name", MODEL_DICT.keys())
+def test_tensorboard(tmp_path, model_name):
+    # Skip if no tensorboard installed
+    pytest.importorskip("tensorboard")
+
+    logname = model_name.upper()
+    algo, env_id = MODEL_DICT[model_name]
+    model = algo('MlpPolicy', env_id, verbose=1, tensorboard_log=tmp_path)
+    model.learn(N_STEPS)
+    model.learn(N_STEPS, reset_num_timesteps=False)
+
+    assert os.path.isdir(tmp_path / str(logname + "_1"))
+    assert not os.path.isdir(tmp_path / str(logname + "_2"))
+
+    logname = "tb_multiple_runs_" + model_name
+    model.learn(N_STEPS, tb_log_name=logname)
+    model.learn(N_STEPS, tb_log_name=logname)
+
+    assert os.path.isdir(tmp_path / str(logname + "_1"))
+    # Check that the log dir name increments correctly
+    assert os.path.isdir(tmp_path / str(logname + "_2"))