[RLlib] Allow MARLModule customization from algorithm config

[kouroshHakha]

Why are these changes needed?

The intent is to allow algorithm level customization of RLModules using RLModuleSpecs to allow maximum flexibility in constructing RLModules (including MARLModules with shared encoders).

Think of allowing users to do this:

class MyCustomRLModuleSpec(SingleRLModuleSpec):
       def build():
              # custom build method

config = config.rl_module(rl_module_spec=MyCustomRLModuleSpec())
algo = config.build()

or

class MyCustomMARLModuleSpec(MultiAgentRLModuleSpec):
       def build():
              shared_encoder = ...
              # custom

config = config.rl_module(rl_module_spec=MyCustomMARLModuleSpec())
algo = config.build()

To achieve this, this PR does a couple of things to enable this:

1. It adds policy_ids to the policy objects, so that they can construct the whole MARLModule but then index the relevant module name in the variable scope of policies and assign the local module to self.model. The reason for this is that we want to allow people to pass in a complicated MARLModule with possible shared encoders, that only are instantiatable via MARLModuleSpecs.
2. The creation of RLModuleSpecs is moved into rollout_workers since there is quite some logic for multi-agent policy_dict and policy map construction inside rollout_worker that allows us to construct MARLModuleSpecs accordingly. If I wanted to write everything from scratch, these would have been created at algorithm level and then passed to all actors (rollout_workers or trainerS), but rollout_worker already does all of that (and also has to work with policies), and for now we leverage it by creating a marl_module_spec inside local_worker and passing a reference in the Algorithm class to pass it to trainer_runner.
3. build_policy_map did a lot of things beside constructing the policy_map. I broke the method into multiple modular stages that should be identical to the original one. The only difference is that when RLModule API is enabled, the marl_module_spec will get constructed by intercepting the policy_specs during build_policy_map operation.

The main changes are in:

rollout_worker.py
algorithm.py
policy.py
algorithm_config.py

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

[gjoliver]

I start to feel like mixing up policy and rl_module stuff may not be the best idea. makes both complicated.
I hope this is temporary and just to accomplish some intermediate goals.

[gjoliver]

is this block temporary? it's just bridging between the RLModule and Policy worlds right?
if so, can we add a TODO/Note?

[kouroshHakha]

Yes, it's temporary until policy co-exists with the RLModule / Learner API. Once we re-write sampler / rollout workers to drop policy, then we won't need this method anymore. instead of creating policy_dicts we will directly create marl_module_specs.

[kouroshHakha]

added TODO/Note.

[gjoliver]

just get rid of the test for now?

[kouroshHakha]

I'm gonna fill out the test since it's relevant to this PR. Basically it tests whether this PR was effective.

[kouroshHakha]

I start to feel like mixing up policy and rl_module stuff may not be the best idea. makes both complicated. I hope this is temporary and just to accomplish some intermediate goals.

I know of no better way to make the transition happen. Otherwise I have to change all 100k lines of RLlib code-base together :)

[kouroshHakha]

These activated some tests that were silently filtered out.

[kouroshHakha]

This is added so that if set_weights() throws an error we catch it. This was surfaced during this PR.