Logging custom arrays with RLlib+Tune

RLlib
3

Thanks, but that’s not really what I’m trying to do.
Actually, you can save entire arrays and TB will put it under the histograms if you use the episode.hist_data in the callback.

But in my case I want to do some post processing and visualization, so I need the arrays saved separately.

After some digging I realized there’s the episode.media property (which was introduced recently) that gets passed to the logger callback under episode_media, i.e.:

def on_episode_end(self, *, worker, base_env, policies, episode, env_index,
                       **kwargs):
        envs = base_env.get_unwrapped()

        episode.media["rollout_arrays"] = [
            env.gat_arr() for env in envs
        ]

and then in the logger CB:

def log_trial_result(self, iteration, trial, result):
    arrays = result['episode_media']['rollout_arrays']
    # save, render, whatevs with arrays
2 Likes