Inconsistency in API Usage for action_space and observation_space

[nobodyPerfecZ]

There is an inconsistency in the API regarding the usage of action_space and observation_space across different environments that inherent from MultiAgentEnv. This inconsistency result in difficulties when initializing networks, as some environments, like Hanabi, follow the MultiAgentEnv API, while others partially follow it or do not follow it at all.

Standard API Implementation (MultiAgentEnv):

multi_agent_env.py

    def observation_space(self, agent: str):
        """Observation space for a given agent."""
        return self.observation_spaces[agent]

    def action_space(self, agent: str):
        """Action space for a given agent."""
        return self.action_spaces[agent]

Affected files:

coin_game.py

    def action_space(self, agent_id=None) -> spaces.Discrete:
        """Action space of the environment."""
        return spaces.Discrete(5)

    def observation_space(self) -> spaces.Box:
        """Observation space of the environment."""
        _shape = (3, 3, 4) if self.cnn else (36,)
        return spaces.Box(low=0, high=1, shape=_shape, dtype=jnp.uint8)

overcooked.py

    def action_space(self, agent_id="") -> spaces.Discrete:
        """Action space of the environment. Agent_id not used since action_space is uniform for all agents"""
        return spaces.Discrete(
            len(self.action_set),
            dtype=jnp.uint32
        )

    def observation_space(self) -> spaces.Box:
        """Observation space of the environment."""
        return spaces.Box(0, 255, self.obs_shape)

storm_2p.py

    def action_space(
        self, agent_id: Union[int, None] = None
    ) -> spaces.Discrete:
        """Action space of the environment."""
        return spaces.Discrete(len(Actions))

    def observation_space(self) -> spaces.Dict:
        """Observation space of the environment."""
        _shape = (
            (OBS_SIZE, OBS_SIZE, len(Items) - 1 + 4)
            if self.cnn
            else (OBS_SIZE**2 * (len(Items) - 1 + 4),)
        )

        return {
            "observation": spaces.Box(
                low=0, high=1, shape=_shape, dtype=jnp.uint8
            ),
            "inventory": spaces.Box(
                low=0,
                high=NUM_COINS,
                shape=NUM_COIN_TYPES + 4,
                dtype=jnp.uint8,
            ),
        }

storm.py

    def action_space(
        self, agent_id: Union[int, None] = None
    ) -> spaces.Discrete:
        """Action space of the environment."""
        return spaces.Discrete(len(Actions))

    def observation_space(self) -> spaces.Dict:
        """Observation space of the environment."""
        _shape_obs = (
            (self.OBS_SIZE, self.OBS_SIZE, (len(Items)-1) + 10)
            if self.cnn
            else (self.OBS_SIZE**2 * ((len(Items)-1) + 10),)
        )

        return spaces.Box(
                low=0, high=1E9, shape=_shape_obs, dtype=jnp.uint8
            ), _shape_obs

[amacrutherford]

Thanks for spotting this, will fix and will write a unit test to catch this in future :)