[RLlib] Fix type hints for original_batches in callbacks.

[XuehaiPan]

Why are these changes needed?

Fix type hints for original_batches in on_postprocess_trajectory.

Related issue number

N/A

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]

what, that's surprising ...
how can we mix these 2 things together?

[XuehaiPan]

ray/rllib/evaluation/collectors/simple_list_collector.py

Lines 824 to 832 in 18c269c

	pre_batches = {}
	for (eps_id, agent_id), collector in self.agent_collectors.items():
	# Build only if there is data and agent is part of given episode.
	if collector.agent_steps == 0 or eps_id != episode_id:
	continue
	pid = self.agent_key_to_policy_id[(eps_id, agent_id)]
	policy = self.policy_map[pid]
	pre_batch = collector.build(policy.view_requirements)
	pre_batches[agent_id] = (policy, pre_batch)

pre_batches is a dict with value type Tuple[Policy, SampleBatch]. Then it is fed to on_postprocess_trajectory as name original_batches.

ray/rllib/evaluation/collectors/simple_list_collector.py

Lines 913 to 925 in 18c269c

	for agent_id, post_batch in sorted(post_batches.items()):
	agent_key = (episode_id, agent_id)
	pid = self.agent_key_to_policy_id[agent_key]
	policy = self.policy_map[pid]
	self.callbacks.on_postprocess_trajectory(
	worker=get_global_worker(),
	episode=episode,
	agent_id=agent_id,
	policy_id=pid,
	policies=self.policy_map,
	postprocessed_batch=post_batch,
	original_batches=pre_batches,
	)

[gjoliver]

oh it's Tuple, not Union, I got scared.
Thanks!

[sven1977]

:) @gjoliver

I think there is only very few places in RLlib where we mix different types e.g. in a return value (for example in the sampler code _process_observations()) and no, we probably shouldn't do this.

[gjoliver]

@sven1977 can you help merge?

[sven1977]

Nice fix. Thanks @XuehaiPan , really appreciate all your help on RLlib :)