[RLlib] PPO torch RLTrainer

[kouroshHakha]

Why are these changes needed?

This PR creates a PPO torch RLTrainer.
PRs that have to merge first:

1. [RLlib] Torch trainer #31628
2. [RLlib] Add algorithm config to RLTrainer #31800
3. [RLlib] Single agent RLTrainer made easy #31802

PRs that make clean ups and are less hacky for this to work:

1. [RLlib] Reparameterize the construction of TrainerRunner and RLTrainers #31991
2. [RLlib] Error out if action_dict is empty in MultiAgentEnv.  #32129
3. [RLlib] Exclude gpu tag from Examples test suite in RLlib #32141
4. [RLlib] Fix revert of trainer runner #32146

Here is the learning curves for CartPole-v1. Blue is the old training stack with RLModules, the red is the new training stack on one CPU. Throughput is also relatively good compared to before.

Let's check it out in multi-gpu case:

Blue is zero gpu, red is 1 gpu, and cyan is 2 gpus.

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

[kouroshHakha]

I have to rebase once the pre-reqs are merged.

[gjoliver]

I have some comments.

[gjoliver]

overriding user configuration is confusing.
maybe consider validate and raise error?

[kouroshHakha]

it just implements what's in the docs, if num_gpus_trainer_worker > 1 and num_trainer_workers = 0 we will just use one gpu. This is just enforcing that from tune's perspective.

[gjoliver]

ok. I just want to provide my personal perspective here.
I feel like we always try to "auto-correct" for our users. for example, compute and override some parameters based on values of other parameters. we also put up a lot of explanations / docs around the "corrections" we may potentially do.
if it's 100% up to me, I would instead just raise error and telling users that I am seeing contradictory configs, because num_trainer_workers=0 doesn't work with num_gpus_per_trainer_worker > 0. now, we have no idea what the actual user intention is. did they mis-specify num_trainer_workers or did they mis-specify num_gpus_per_trainer_worker? the only one who can fix this is the user him/herself.
now, this is just some of my thoughts. I can't enforce this, so hopefully you understand what I meant.

[kouroshHakha]

got it. make sense.

[kouroshHakha]

done

[gjoliver]

what does remote .use_critic mean?

[kouroshHakha]

typo: remove :) not remote

[gjoliver]

I actually think you should be explicit, and not rely on some keys on the result dict to tell you which policies need to be updated.
train_results should not be used as control messages basically.

[kouroshHakha]

I'd like to hear more about what you mean by being more explicit? I was planning on revisiting the train_results structure to remove these requirements in the next round of updates, But would love to hear your thoughts on how it should ideally look like?

[gjoliver]

I am not opinionated about how the result dict should look like.
but I do think we shouldn't use it as a control message, meaning, I can only get my policies updated if I add something in the result dict.
these two things probably shouldn't go together?

[kouroshHakha]

I see your point and I actually have a better idea? right now I haven't even made trainer_runner such that it only updates the policies that is allowed (via policies_to_train). I think with that variable lingering around, I can infer the policies to update. Then the retuened results won't be used as the message passing medium. I'll add a todo with better design guideline in the next PR.

[gjoliver]

you duplicate module_batches multiple times to make a mini_batch large enough?
this block of code is actually very confusing ...

[kouroshHakha]

yes, generally speaking, if number of samples within each policy batch is un-even and skewed, you need to make sure that they are sharded almost equally when you pass them to the RLTrainer. Say one policy1 have 100 samples, and policy2 has 20 samples. but when I want to pass a minibatch size of 40, from policy1 I will select 0-39, but from policy2 I select 0-19+0-19 to make up 40 samples. This is the problem of sharding across policies, I couldn't figure out an easy-to-understand sharding strategy. But I want to think about this more when I think about the double batch communication overhead, they are very relevant.

[gjoliver]

I see. maybe make a util function out of this. so then it's pluggable. and we can write and compare a few different schemes.

[gjoliver]

btw, not sure if you want to add these as comments in the code. that would help a lot.

[kouroshHakha]

will update taking into account the suggestions 👍

[kouroshHakha]

ok I am creating a utility MiniBatchIterator that can be re-used once we move this iteration to sharded batches inside RLTrainer as well. Thanks for the suggestion. I really like the design break down.

[gjoliver]

👍