[RLlib] Add 'single_action_space' and 'single_observation_space' to 'VectorMultiAgentEnv`. (#51096)

Signed-off-by: Abrar Sheikh <abrar@anyscale.com>

Author: simonsays1980

rllib/env/vector/sync_vector_multi_agent_env.py

@@ -37,10 +37,11 @@ def __init__(
         self.single_action_spaces = self.envs[0].unwrapped.action_spaces or dict(
             self.envs[0].unwrapped.action_space
         )
+        self.single_action_space = gym.spaces.Dict(self.single_action_spaces)
         self.single_observation_spaces = self.envs[
             0
         ].unwrapped.observation_spaces or dict(self.envs[0].unwrapped.observation_space)
-
+        self.single_observation_space = gym.spaces.Dict(self.single_observation_spaces)
         # TODO (simon): Decide if we want to include a spaces check here.
 
         # Note, if `single_observation_spaces` are not defined, this will

rllib/env/vector/vector_multi_agent_env.py

@@ -25,6 +25,9 @@ class VectorMultiAgentEnv:
     # single on that holds for all sub-envs.
     single_observation_spaces: Optional[Dict[str, gym.Space]] = None
     single_action_spaces: Optional[Dict[str, gym.Space]] = None
+    # Note, the proper `gym` spaces are needed for the connector pipeline.
+    single_observation_space: Optional[gym.spaces.Dict] = None
+    single_action_space: Optional[gym.spaces.Dict] = None
 
     num_envs: int