From 5993033c73d56e8ac13f59647d47374766d0b63d Mon Sep 17 00:00:00 2001 From: "Carlos M. Casas Cuadrado" Date: Fri, 8 Jan 2021 15:47:08 +0100 Subject: [PATCH] Add image and figure to tensorboard logger (#277) * Added Image and Figure classes to logger. For now, these objects can only be logged by TensorBoardOutputFormat * Added documentation for figure and image logging into tensorboard * Updated changelog * Minor changes to documentation. Reviewed supported types for logging images and figures * Fix type for np arrays * Added more explicit example for logging figures in the documentation. Added docstrings for parameters in logging auxiliary classes * Added tests for image and figure logging * Applied autoformatting * Update doc * Fix documentation example * Bump version Co-authored-by: Carlos Casas Co-authored-by: Antonin RAFFIN --- docs/guide/tensorboard.rst | 72 ++++++++++++++++++++++++++++++ docs/misc/changelog.rst | 5 ++- stable_baselines3/common/logger.py | 54 ++++++++++++++++++++++ stable_baselines3/version.txt | 2 +- tests/test_logger.py | 51 +++++++++++++++++++++ 5 files changed, 181 insertions(+), 3 deletions(-) diff --git a/docs/guide/tensorboard.rst b/docs/guide/tensorboard.rst index cf5c4cf..788ccc1 100644 --- a/docs/guide/tensorboard.rst +++ b/docs/guide/tensorboard.rst @@ -81,6 +81,78 @@ Here is a simple example on how to log both additional tensor or arbitrary scala model.learn(50000, callback=TensorboardCallback()) +Logging Images +-------------- + +TensorBoard supports periodic logging of image data, which helps evaluating agents at various stages during training. + +.. warning:: + To support image logging `pillow `_ must be installed otherwise, TensorBoard ignores the image and logs a warning. + +Here is an example of how to render an image to TensorBoard at regular intervals: + +.. code-block:: python + + from stable_baselines3 import SAC + from stable_baselines3.common.callbacks import BaseCallback + from stable_baselines3.common.logger import Image + + model = SAC("MlpPolicy", "Pendulum-v0", tensorboard_log="/tmp/sac/", verbose=1) + + + class ImageRecorderCallback(BaseCallback): + def __init__(self, verbose=0): + super(ImageRecorderCallback, self).__init__(verbose) + + def _on_step(self): + image = self.training_env.render(mode="rgb_array") + # "HWC" specify the dataformat of the image, here channel last + # (H for height, W for width, C for channel) + # See https://pytorch.org/docs/stable/tensorboard.html + # for supported formats + self.logger.record("trajectory/image", Image(image, "HWC"), exclude=("stdout", "log", "json", "csv")) + return True + + + model.learn(50000, callback=ImageRecorderCallback()) + +Logging Figures/Plots +--------------------- +TensorBoard supports periodic logging of figures/plots created with matplotlib, which helps evaluating agents at various stages during training. + +.. warning:: + To support figure logging `matplotlib `_ must be installed otherwise, TensorBoard ignores the figure and logs a warning. + +Here is an example of how to store a plot in TensorBoard at regular intervals: + +.. code-block:: python + + import numpy as np + import matplotlib.pyplot as plt + + from stable_baselines3 import SAC + from stable_baselines3.common.callbacks import BaseCallback + from stable_baselines3.common.logger import Figure + + model = SAC("MlpPolicy", "Pendulum-v0", tensorboard_log="/tmp/sac/", verbose=1) + + + class FigureRecorderCallback(BaseCallback): + def __init__(self, verbose=0): + super(FigureRecorderCallback, self).__init__(verbose) + + def _on_step(self): + # Plot values (here a random variable) + figure = plt.figure() + figure.add_subplot().plot(np.random.random(3)) + # Close the figure after logging it + self.logger.record("trajectory/figure", Figure(figure, close=True), exclude=("stdout", "log", "json", "csv")) + plt.close() + return True + + + model.learn(50000, callback=FigureRecorderCallback()) + Logging Videos -------------- diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 726c9ad..401f414 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Pre-Release 0.11.0a4 (WIP) +Pre-Release 0.11.0a5 (WIP) ------------------------------- Breaking Changes: @@ -26,6 +26,7 @@ New Features: - Added ``monitor_kwargs`` parameter to ``make_vec_env`` and ``make_atari_env`` - Wrap the environments automatically with a ``Monitor`` wrapper when possible. - ``EvalCallback`` now logs the success rate when available (``is_success`` must be present in the info dict) +- Added new wrappers to log images and matplotlib figures to tensorboard. (@zampanteymedio) Bug Fixes: ^^^^^^^^^^ @@ -534,4 +535,4 @@ And all the contributors: @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray -@tfederico @hn2 @LucasAlegre @AptX395 +@tfederico @hn2 @LucasAlegre @AptX395 @zampanteymedio diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py index 391c62e..10283ec 100644 --- a/stable_baselines3/common/logger.py +++ b/stable_baselines3/common/logger.py @@ -10,6 +10,7 @@ from typing import Any, Dict, List, Optional, Sequence, TextIO, Tuple, Union import numpy as np import pandas import torch as th +from matplotlib import pyplot as plt try: from torch.utils.tensorboard import SummaryWriter @@ -26,6 +27,9 @@ DISABLED = 50 class Video(object): """ Video data class storing the video frames and the frame per seconds + + :param frames: frames to create the video from + :param fps: frames per second """ def __init__(self, frames: th.Tensor, fps: Union[float, int]): @@ -33,6 +37,34 @@ class Video(object): self.fps = fps +class Figure(object): + """ + Figure data class storing a matplotlib figure and whether to close the figure after logging it + + :param figure: figure to log + :param close: if true, close the figure after logging it + """ + + def __init__(self, figure: plt.figure, close: bool): + self.figure = figure + self.close = close + + +class Image(object): + """ + Image data class storing an image and data format + + :param image: image to log + :param dataformats: Image data format specification of the form NCHW, NHWC, CHW, HWC, HW, WH, etc. + More info in add_image method doc at https://pytorch.org/docs/stable/tensorboard.html + Gym envs normally use 'HWC' (channel last) + """ + + def __init__(self, image: Union[th.Tensor, np.ndarray, str], dataformats: str): + self.image = image + self.dataformats = dataformats + + class FormatUnsupportedError(NotImplementedError): def __init__(self, unsupported_formats: Sequence[str], value_description: str): if len(unsupported_formats) > 1: @@ -108,6 +140,12 @@ class HumanOutputFormat(KVWriter, SeqWriter): if isinstance(value, Video): raise FormatUnsupportedError(["stdout", "log"], "video") + if isinstance(value, Figure): + raise FormatUnsupportedError(["stdout", "log"], "figure") + + if isinstance(value, Image): + raise FormatUnsupportedError(["stdout", "log"], "image") + if isinstance(value, float): # Align left value_str = f"{value:<8.3g}" @@ -196,6 +234,10 @@ class JSONOutputFormat(KVWriter): def cast_to_json_serializable(value: Any): if isinstance(value, Video): raise FormatUnsupportedError(["json"], "video") + if isinstance(value, Figure): + raise FormatUnsupportedError(["json"], "figure") + if isinstance(value, Image): + raise FormatUnsupportedError(["json"], "image") 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 @@ -258,6 +300,12 @@ class CSVOutputFormat(KVWriter): if isinstance(value, Video): raise FormatUnsupportedError(["csv"], "video") + if isinstance(value, Figure): + raise FormatUnsupportedError(["csv"], "figure") + + if isinstance(value, Image): + raise FormatUnsupportedError(["csv"], "image") + if value is not None: self.file.write(str(value)) self.file.write("\n") @@ -296,6 +344,12 @@ class TensorBoardOutputFormat(KVWriter): if isinstance(value, Video): self.writer.add_video(key, value.frames, step, value.fps) + if isinstance(value, Figure): + self.writer.add_figure(key, value.figure, step, close=value.close) + + if isinstance(value, Image): + self.writer.add_image(key, value.image, step, dataformats=value.dataformats) + # Flush the output to the file self.writer.flush() diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 1b742ef..f23fba9 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.11.0a4 +0.11.0a5 diff --git a/tests/test_logger.py b/tests/test_logger.py index 98f832d..c1cce85 100644 --- a/tests/test_logger.py +++ b/tests/test_logger.py @@ -3,11 +3,14 @@ from typing import Sequence import numpy as np import pytest import torch as th +from matplotlib import pyplot as plt from pandas.errors import EmptyDataError from stable_baselines3.common.logger import ( DEBUG, + Figure, FormatUnsupportedError, + Image, ScopedConfigure, Video, configure, @@ -198,3 +201,51 @@ def test_report_video_to_unsupported_format_raises_error(tmp_path, unsupported_f writer.write({"video": video}, key_excluded={"video": ()}) assert unsupported_format in str(exec_info.value) writer.close() + + +def test_report_image_to_tensorboard(tmp_path, read_log): + pytest.importorskip("tensorboard") + + image = Image(image=th.rand(16, 16, 3), dataformats="HWC") + writer = make_output_format("tensorboard", tmp_path) + writer.write({"image": image}, key_excluded={"image": ()}) + + assert not read_log("tensorboard").empty + writer.close() + + +@pytest.mark.parametrize("unsupported_format", ["stdout", "log", "json", "csv"]) +def test_report_image_to_unsupported_format_raises_error(tmp_path, unsupported_format): + writer = make_output_format(unsupported_format, tmp_path) + + with pytest.raises(FormatUnsupportedError) as exec_info: + image = Image(image=th.rand(16, 16, 3), dataformats="HWC") + writer.write({"image": image}, key_excluded={"image": ()}) + assert unsupported_format in str(exec_info.value) + writer.close() + + +def test_report_figure_to_tensorboard(tmp_path, read_log): + pytest.importorskip("tensorboard") + + fig = plt.figure() + fig.add_subplot().plot(np.random.random(3)) + figure = Figure(figure=fig, close=True) + writer = make_output_format("tensorboard", tmp_path) + writer.write({"figure": figure}, key_excluded={"figure": ()}) + + assert not read_log("tensorboard").empty + writer.close() + + +@pytest.mark.parametrize("unsupported_format", ["stdout", "log", "json", "csv"]) +def test_report_figure_to_unsupported_format_raises_error(tmp_path, unsupported_format): + writer = make_output_format(unsupported_format, tmp_path) + + with pytest.raises(FormatUnsupportedError) as exec_info: + fig = plt.figure() + fig.add_subplot().plot(np.random.random(3)) + figure = Figure(figure=fig, close=True) + writer.write({"figure": figure}, key_excluded={"figure": ()}) + assert unsupported_format in str(exec_info.value) + writer.close()