Added support for stateful agents.

Signed-off-by: Matthew <[email protected]>

Author: MatthewCWeston

rllib/examples/algorithms/mappo/connectors/general_advantage_estimation.py

@@ -64,11 +64,11 @@ def __call__(
             and (not isinstance(rl_module[k], SelfSupervisedLossAPI))
         ]
         critic_batch[Columns.OBS] = torch.cat(
-            [batch[k][Columns.OBS] for k in obs_mids], dim=1
+            [batch[k][Columns.OBS] for k in obs_mids], dim=-1
         )
         # Compute value predictions
         vf_preds = rl_module[SHARED_CRITIC_ID].compute_values(critic_batch)
-        vf_preds = {mid: vf_preds[:, i] for i, mid in enumerate(obs_mids)}
+        vf_preds = {mid: vf_preds[..., i] for i, mid in enumerate(obs_mids)}
         # Loop through all modules and perform each one's GAE computation.
         for module_id, module_vf_preds in vf_preds.items():
             module = rl_module[module_id]
@@ -136,10 +136,10 @@ def __call__(
             batch[module_id][Postprocessing.VALUE_TARGETS] = module_value_targets
         # Add GAE results to the critic batch
         critic_batch[Postprocessing.VALUE_TARGETS] = np.stack(
-            [batch[mid][Postprocessing.VALUE_TARGETS] for mid in obs_mids], axis=1
+            [batch[mid][Postprocessing.VALUE_TARGETS] for mid in obs_mids], axis=-1
         )
         critic_batch[Postprocessing.ADVANTAGES] = np.stack(
-            [batch[mid][Postprocessing.ADVANTAGES] for mid in obs_mids], axis=1
+            [batch[mid][Postprocessing.ADVANTAGES] for mid in obs_mids], axis=-1
         )
         batch[SHARED_CRITIC_ID] = critic_batch  # Critic data -> training batch
         # Convert all GAE results to tensors.

rllib/examples/multi_agent/pettingzoo_shared_value_function.py

@@ -47,6 +47,7 @@
 
 from pettingzoo.sisl import waterworld_v4
 
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
 from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
 from ray.rllib.core.rl_module.rl_module import RLModuleSpec
 from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv
@@ -68,6 +69,11 @@
     default_timesteps=1000000,
     default_reward=0.0,
 )
+parser.add_argument(
+    "--use-lstm",
+    action="store_true",
+    help="Whether to use LSTM encoders for the agents' policies.",
+)
 
 
 if __name__ == "__main__":
@@ -87,7 +93,10 @@ def get_env(_):
 
     # An agent for each of our policies, and a single shared critic
     env_instantiated = get_env({})  # neccessary for non-agent modules
-    specs = {p: RLModuleSpec() for p in policies}
+    model_config = DefaultModelConfig(
+        use_lstm=args.use_lstm,
+    )
+    specs = {p: RLModuleSpec(model_config=model_config) for p in policies}
     specs[SHARED_CRITIC_ID] = RLModuleSpec(
         module_class=SharedCriticTorchRLModule,
         observation_space=env_instantiated.observation_space[policies[0]],