diff --git a/supersuit/multiagent_wrappers/black_death.py b/supersuit/multiagent_wrappers/black_death.py index 0cc06c2c..0bd472b3 100644 --- a/supersuit/multiagent_wrappers/black_death.py +++ b/supersuit/multiagent_wrappers/black_death.py @@ -5,16 +5,35 @@ from supersuit.utils.wrapper_chooser import WrapperChooser +def _assert_black_deathable(space): + """Black death needs a well-defined zero observation. + + Box has one. Dict has one exactly when all of its subspaces do, so it is + checked recursively. + """ + if isinstance(space, gymnasium.spaces.Dict): + for subspace in space.spaces.values(): + _assert_black_deathable(subspace) + else: + assert isinstance( + space, gymnasium.spaces.Box + ), f"observation spaces for black death must be Box or Dict spaces, is {space}" + + +def _zero_obs(space): + """Zero observation matching the structure of space.""" + if isinstance(space, gymnasium.spaces.Dict): + return {name: _zero_obs(subspace) for name, subspace in space.spaces.items()} + return np.zeros_like(space.low) + + class black_death_par(BaseParallelWrapper): def __init__(self, env): super().__init__(env) def _check_valid_for_black_death(self): for agent in self.agents: - space = self.observation_space(agent) - assert isinstance( - space, gymnasium.spaces.Box - ), f"observation sapces for black death must be Box spaces, is {space}" + _assert_black_deathable(self.observation_space(agent)) def reset(self, seed=None, options=None): obss, infos = self.env.reset(seed=seed, options=options) @@ -22,7 +41,7 @@ def reset(self, seed=None, options=None): self.agents = self.env.agents[:] self._check_valid_for_black_death() black_obs = { - agent: np.zeros_like(self.observation_space(agent).low) + agent: _zero_obs(self.observation_space(agent)) for agent in self.agents if agent not in obss } @@ -32,7 +51,7 @@ def step(self, actions): active_actions = {agent: actions[agent] for agent in self.env.agents} obss, rews, terms, truncs, infos = self.env.step(active_actions) black_obs = { - agent: np.zeros_like(self.observation_space(agent).low) + agent: _zero_obs(self.observation_space(agent)) for agent in self.agents if agent not in obss } diff --git a/test/parallel_env_test.py b/test/parallel_env_test.py index d0c5dc0c..5d2456ba 100644 --- a/test/parallel_env_test.py +++ b/test/parallel_env_test.py @@ -1,6 +1,6 @@ import numpy as np import pytest -from gymnasium.spaces import Box, Discrete +from gymnasium.spaces import Box, Dict, Discrete from pettingzoo.utils import ParallelEnv import supersuit @@ -58,6 +58,23 @@ def close(self): } base_act_spaces = {f"a{idx}": Discrete(5) for idx in range(2)} +dict_obs = { + f"a{idx}": { + "camera": np.zeros([8, 8, 3], dtype=np.float32) + idx, + "position": np.zeros([3], dtype=np.float32) + idx, + } + for idx in range(2) +} +dict_obs_space = { + f"a{idx}": Dict( + { + "camera": Box(low=np.float32(0.0), high=np.float32(10.0), shape=[8, 8, 3]), + "position": Box(low=np.float32(-1.0), high=np.float32(1.0), shape=[3]), + } + ) + for idx in range(2) +} + def test_basic(): env = DummyParEnv(base_obs, base_obs_space, base_act_spaces) @@ -90,3 +107,59 @@ def test_black_death_done_semantics(terms_val, truncs_val, expected_done): assert all(terms.values()) == expected_done assert all(truncs.values()) == expected_done assert (len(env.agents) == 0) == expected_done + + +def test_black_death_dict_obs(): + """black_death must accept Dict observation spaces and zero-fill per key.""" + env = DummyParEnv(dict_obs, dict_obs_space, base_act_spaces) + env = supersuit.black_death_v3(env) + obs, _ = env.reset() + + assert set(obs) == {"a0", "a1"} + for agent_obs in obs.values(): + assert set(agent_obs) == {"camera", "position"} + assert agent_obs["camera"].shape == (8, 8, 3) + assert agent_obs["position"].shape == (3,) + + actions = {agent: env.action_space(agent).sample() for agent in env.agents} + obs, _, _, _, _ = env.step(actions) + for agent_obs in obs.values(): + assert agent_obs["camera"].shape == (8, 8, 3) + + +def test_black_death_dict_obs_zero_fills_dropped_agent(): + """An agent missing from the step output gets a structured zero observation.""" + env = DummyParEnv(dict_obs, dict_obs_space, base_act_spaces) + wrapped = supersuit.black_death_v3(env) + wrapped.reset() + + # drop a1 from the underlying env's output + env._observations = {"a0": dict_obs["a0"]} + env.agents = ["a0"] + env.rewards = {"a0": 1} + env.terminations = {"a0": False} + env.truncations = {"a0": False} + env.infos = {"a0": {}} + + obs, _, _, _, _ = wrapped.step({"a0": wrapped.action_space("a0").sample()}) + + assert set(obs) == {"a0", "a1"} + assert np.all(obs["a1"]["camera"] == 0) + assert np.all(obs["a1"]["position"] == 0) + assert obs["a1"]["camera"].shape == (8, 8, 3) + assert obs["a1"]["position"].shape == (3,) + + +def test_black_death_rejects_unsupported_subspace(): + from gymnasium.spaces import MultiBinary + + bad_spaces = { + f"a{idx}": Dict( + {"ok": Box(low=0.0, high=1.0, shape=[3]), "bad": MultiBinary(4)} + ) + for idx in range(2) + } + env = DummyParEnv(dict_obs, bad_spaces, base_act_spaces) + env = supersuit.black_death_v3(env) + with pytest.raises(AssertionError): + env.reset()