[feature request] Support Tensorboard video logging (convenience)

When investigating RL algorithms I’ve found it quite useful to report videos of a trajectory to Tensorboard, to directly relate it to other performance measures and look at the agent at different stages.

The implementation would be quite simple I think and if you decide this to be a useful feature I’d be happy to implement it. As a matter of fact I’ve already done so in my own fork of this repo. Here is how I envisioned its usage:

def render_trajectory(env, model):
    obs = env.reset()
    rewards = []
    done = False
    screens = []
    while not done:
        actions, _ = model.predict(obs, deterministic=True)
        obs, _, done, _ = env.step(actions)
        screen = env.render(mode='rgb_array')
        screens.append(screen.transpose(2, 0, 1))

    return Video(th.ByteTensor([screens]), fps=40)

class ReportTrajectoryCallback(BaseCallback):
    def __init__(self, eval_env, check_freq):
        super().__init__()
        self._eval_env = eval_env
        self._check_freq = check_freq

    def _on_step(self):
        if self.n_calls % self._check_freq == 0:
            video = render_trajectory(self._eval_env, self.model)
            self.logger.record("trajectory/video", video)
        return True

env = gym.make('CartPole-v1')
model = A2C('MlpPolicy', env)
model.learn(total_timesteps=int(1e4), callback=ReportTrajectoryCallback(env, 1000))

I’d suggest introducing a Video object and a simple extension of the logger to record the video. The TensorBoardOutputFormat would simply check for a video object and forward it to the SummaryWriter. All other KVWriter classes would just ignore it. As mentioned I’ve done this already for me self. Here is the relevant commit: https://github.com/SwamyDev/stable-baselines3/commit/d6b15416000011a5aeddd4af1bfb9f9a2d79d22f

Of course I’m open to any different approaches. For instance checking the type directly in the Formater classes is not good object oriented style and maybe an explicit video logging function would be better. However, it is relatively simple and get’s the job done.

What do you think? Would this is be a small but useful contribution?

[araffin]

In fact, part of it should be included inside SB3 repo, because we need to change the tensorboard writer. But I agree that the callback and the helper should go in the contrib folder.

And there is a simpler way to use it: record("trajectory/video", video, exclude=("json", "stdout", "log", "csv"))

[SwamyDev]

Ah thanks for the suggestion that makes my code a lot simpler 😃

Ok, I’ll file a proper bug report and everything for the other issue