docstr: Tests creating a customized MARL BC module that owns a shared encoder.
docstr_tok: ["Tests","creating","a","customized","MARL","BC","module","that","owns","a","shared","encoder","."]
code: def test_customized_multi_agent_module(self):
      """Tests creating a customized MARL BC module that owns a shared encoder."""
  
      global_dim = 10
      local_dims = [16, 32]
      action_dims = [2, 4]
  
      # TODO (Kourosh): add tf support
      for fw in ["torch"]:
          marl_module_cls = CUSTOM_MARL_MODULES[fw]
          rl_module_cls = CUSTOM_MODULES[fw]
  
          spec = MultiAgentRLModuleSpec(
              marl_module_class=marl_module_cls,
              module_specs={
                  "agent_1": SingleAgentRLModuleSpec(
                      module_class=rl_module_cls,
                      observation_space=gym.spaces.Dict(
                          {
                              "global": gym.spaces.Box(
                                  low=-1, high=1, shape=(global_dim,)
                              ),
                              "local": gym.spaces.Box(
                                  low=-1, high=1, shape=(local_dims[0],)
                              ),
                          }
                      ),
                      action_space=gym.spaces.Discrete(action_dims[0]),
                      model_config_dict={"fcnet_hiddens": [128]},
                  ),
                  "agent_2": SingleAgentRLModuleSpec(
                      module_class=rl_module_cls,
                      observation_space=gym.spaces.Dict(
                          {
                              "global": gym.spaces.Box(
                                  low=-1, high=1, shape=(global_dim,)
                              ),
                              "local": gym.spaces.Box(
                                  low=-1, high=1, shape=(local_dims[1],)
                              ),
                          }
                      ),
                      action_space=gym.spaces.Discrete(action_dims[1]),
                      model_config_dict={"fcnet_hiddens": [128]},
                  ),
              },
          )
  
          model = spec.build()
  
          if fw == "torch":
              # change the parameters of the shared encoder and make sure it changes
              # across all agents
              foo = model["agent_1"].encoder[0].bias
              foo.data = torch.ones_like(foo.data)
              self.assertTrue(
                  torch.allclose(model["agent_2"].encoder[0].bias, foo)
              )
  
file_path: ray/rllib/core/rl_module/tests/test_rl_module_specs.py
imports: import unittest
  import gymnasium
  import torch
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec
  from ray.rllib.core.rl_module.marl_module import MultiAgentRLModule, MultiAgentRLModuleSpec
  from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule, BCTorchRLModuleWithSharedGlobalEncoder, BCTorchMultiAgentModuleWithSharedEncoder
  from ray.rllib.core.testing.tf.bc_module import DiscreteBCTFModule, BCTfRLModuleWithSharedGlobalEncoder, BCTfMultiAgentModuleWithSharedEncoder
label_desc: Use image node_id 3 for calling the TestRLModuleSpecs obj's underlying member method code with example usage: obj.test_customized_multi_agent_module() without return types