[RLlib] Added functionality to add infos and extra_model_outputs to the sample output of PrioritizedEpisodeReplayBuffer.

[simonsays1980]

Why are these changes needed?

So far PrioritizedEpisodeReplayBuffer had a functionality to add infos to the sample of this buffer, but not one to add also extra_model_outputs. This PR adds the functionality together with a corresponding test case.

Note, the extra_model_outputs are extracted as a dict and will be added to the batch in this form per row (similar to infos). Later in post-processing the variables from this dicitonary can be extracted in a corresponding learner connector. Furthermore, while infos are extracted at the end of n_step, the extra_model_outputs usually refer to a corresponding action which comes from the first timestep in the n_step tuple. Henceforth, we take the extra_model_outputs from the same timestep.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• 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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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]

Not sure this is a good idea just np'ing stuff like this. This often leads to these unwieldy object arrays that have unpredictable behavior (the same is true for np'ing the infos above, we should just keep them as a list of infos-dicts in the returned batch).

We usually separate these sub-columns in extra_model_outputs in our batches. Can we do that here, too?

ret.update(
{
    k: batch(v)
    for k, v in extra_model_outputs.items()
}
)

The final batch (returned from sample) should have columns at the top level, e.g. OBS or ACTION_DIST_INPUTS.
Under each of these columns should be a (possibly nested) struct of numpy array leafs (or simply a numpy array if no complex space/struct). All leafs should have the shape (B, T?, ...), where T might be 0 or 1.

Let me know, if I'm making a thinking-mistake here. :)

[simonsays1980]

Hey @sven1977 thanks for the review! Yes this was somehow still ambiguous how to deal with the extra model outputs. I can batch the items from this field such that each of the keys in extra_model_outputs defines a new column in the batch.

[simonsays1980]

@sven1977 following your logic above it might also make sense to keep the other "batch" columns here as lists such that they can be batched in a standard way in the connectors?

[sven1977]

LGTM! Let's merge this, then time-it, whether the saved time to create the batch in the buffer is eaten up by the additional batching step required in the Learner Connector (I don't think that would be the case).

Awesome PR @simonsays1980 ! :)