How to log Render to tensorboard?

Hi there,

I’m trying to log the render of my environment to tensorboard.

Similar to @smorad proposal in tensorboard render log i created the following callback

class VideoCallback(DefaultCallbacks):
    def on_episode_step(self,*,worker,base_env,episode,env_index,**kwargs):
        # render env
        img = base_env.get_unwrapped()[0].render(mode = 'rgb_array')

        if episode.user_data.get('video',False):
            episode.user_data["video"].append(img)
        else:
            episode.user_data["video"] = [img]

    def on_episode_end(self,*,worker, base_env, policies, episode, env_index, **kwargs):
        
        video = np.stack(episode.user_data["video"])
        

        # save to file
        imgs = [Image.fromarray(img) for img in episode.user_data["video"]]
        imgs[0].save("/home/ray/ray_video/render.gif", save_all=True, append_images=imgs[1:], duration=500, loop=0)

        # expand dimension for tf logger
        video = np.expand_dims(video,axis=0)
        
        # add to metric
        episode.custom_metrics["my_video"] = video
        episode.user_data["video"]=None

        # print('episode running', episode, env_index)
        # print(video.shape)

saving the render as gif works fine. My question is: How do I correctly add a video to the Tensorboard Logger?

At the moment I tried saving the rgb array as a custom metric. This fails, because i submit an array and not a number and .isnan returns a boolean array in ray/metrics.py at e5feaee95a6c3fa37a7b59318aebb49940f19617 · ray-project/ray · GitHub

I suspect I understood something wrong in how one uses the custom metrics for this. Am I correct or do i have to change the evaluation from

for k, v_list in custom_metrics.copy().items():
        filt = [v for v in v_list if not np.isnan(v)]

to

for k, v_list in custom_metrics.copy().items():
        filt = [v for v in v_list if not np.any(np.isnan(v))]

and are there other problems with my approach?

You are on the right path. I admit, video recording in TB is not so well documented/supported. Could you try adding/editing the elif-code block like this instead?

from ray/python/ray/tune/logger.py::~206

            elif ((isinstance(value, list) and len(value) > 0)
                  or (isinstance(value, np.ndarray) and value.size > 0)):
                valid_result[full_attr] = value

                # Single video or list of videos.
                if isinstance(value, np.ndarray) and value.ndim == 5:
                    self._file_writer.add_video(
                        full_attr, value, global_step=step, fps=20)
                    continue
                elif isinstance(value, list) and isinstance(value[0], np.ndarray) \
                        and value[0].ndim == 5:
                    video = np.concatenate(value, axis=1)
                    self._file_writer.add_video(
                        full_attr, video, global_step=step, fps=20)
                    continue

                try:
                   ...

It’s also saying here: tensorboardX — tensorboardX documentation

that you need to do pip install moviepy for this to work.

When I do this and run a simple CartPole example with your CallBacks class, I can click on “Images” in TB, but I don’t see any videos there, just a seemingly empty panel. I’ll keep trying to get this to work reliably.

Oh, also, you have to specify your logger type in tune.run(..., loggers=[TBXLogger]) with from ray.tune.logger import TBXLogger

1 Like

Thank you very much for your help.

I notices that in ray/rllib/evaluation/metrics.py

        del custom_metrics[k]

The custom metric is deleted. Commenting this line is workaround.

For the cases where one can not see any render:
make sure the dimensions are correct. The tensorboardX logger expects:

    #Expected (N,T,C=3,H,W)
    video = np.moveaxis(video,3,1)  # for most envs

I hope this is useful also to others reading it.