[RLlib] Example script custom_metrics_and_callbacks.py should work for batch_mode=complete_episodes.

[simonsays1980]

Why are these changes needed?

To enable a better traceability and to give users a universal guideline that works with different batch_modes when producing custom metrics.
So far the batch_mode="complete_episodes" produces an error in the custom_metrics_and_callbacks.py. The reason for this is:

• When batch_mode="complete_episodes" the SimpleListCollector will be called to build a MultiAgentBatch
• That in turn will call the build() method of the PolicyCollector and empty the batches attribute.
• As the batches from the PolicyCollector are empty now, episode.batch_builder.policy_collectors["default_policy"].batches[-1] does not exist in the assert expression.

For most of the end-users running their algorithms with batch_mode="complete_episodes" trying to implement their own metrics by following this example will possibly not able to trace this error back and to find a solution for their own code.

Related issue number

#22683

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 :(

[gjoliver]

can you please rebase master into your branch, a lot of strange numpy random number generator errors.
thanks.

[gjoliver]

thanks a ton for noticing this.
any suggestion on how can we assert this for "complete_episode" mode as well?

[simonsays1980]

Trainer class

I think in case of complete_episodes we do not have to care about, if episodes are really done because they are always done (complete). But to have a callback that can run with both batch_modes we need a prior safeguard for this. In case of truncate_episodes we have to be careful if the episode is indeed done.

Custom usage

This is different, if a user implements the callback in his own code and calls this callback before an episode is done. I don't know if (1) this case is rare and would happen only between a rollout and a batch completion or (2) if it can happen literally at any point of sampling. If we want to ensure that the below assertion does not error out, we might have a look at the batches attribute of the PolicyCollector as this gets emptied when the MultiAgentBatch is created by the simple_list_collector. I have also not yet figured out, if this attribute (batches) is only emptied, when the episode is done.

[sven1977]

Actually, even in "complete_episodes" mode, there could be more than one episodes in the final train batch (but not here in this callback).

• We should probably rename multiple_episodes_in_batch into collect_complete_episodes (bool).
• I think we should rather use worker.policy_config["batch_mode"] here in this if block.

[sven1977]

Hey @simonsays1980 , thanks for this fix! Let me know, if you want to do a follow-up PR with the 2 above suggested changes. I think this would help clarify these batch generating rules for the different settings even more.

[simonsays1980]

Hi @sven1977, yes I can do a follow-up PR. Should I do this within the same branch or better create a new one?

[simonsays1980]

@gjoliver I rebased. Can we rerun the tests somehow?

[gjoliver]

ok, the tests look reasonable.
@sven1977, can you help merge? :)