[RLLib] Episode media logging support

[Bam4d]

Why are these changes needed?

There is currently no way of transporting episode data such as images or videos for custom environments such as multi-agent environments to logging callbacks such as WandbLoggerCallback.

This patch allows agent callbacks to add arbitrary serializable media data such as wandb.Video() to an episode.media object which is then collated and returned as part of the result (which is used in both training and evaluation)

result = trainer.train()
my_custom_media = result['episode_media']

The episode media can be customized in any of the episode callbacks, for example on episode end:

    def on_episode_end(self, *, worker: "RolloutWorker", base_env: BaseEnv, policies: Dict[PolicyID, Policy],
                       episode: MultiAgentEpisode, env_index: Optional[int] = None, **kwargs) -> None:

        if not worker.multiagent:
            info = episode.last_info_for()
            if 'video' in info:
                level = info['video']['level']
                path = info['video']['path']
                print(f'creating video with path: {path}')
                episode.media[f'level_{level}'] = Video(path)

In summary this provides generic support for logging any custom environment media including seperate videos per-agent in mutli-agent enviroments.

Related issue number

#13169
#8403

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
• [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

@sven1977 This is my first (of hopefully many) commits to rllib. It would be great to get some advice on whats expected other than these tick-boxes here. for example I could definitely include a documentation change here, but seems like it makes more sense to get the code reviewed first before writing docs for it.

[jsuarez5341]

Related issue: #12833. This is a big step in the right direction (though the callback api is still a bit limited)

[sven1977]

Looks great! Thanks for providing this generic solution.
We should probably add an example script that demonstrates this with on a simple, custom (maybe multi-agent) env.

[Bam4d]

We should probably add an example script that demonstrates this with on a simple, custom (maybe multi-agent) env.

Where would the best place for this be? in the docs somewhere?

[rfali]

an example script for this would be great. waiting for it