Support Dict observation spaces in black_death

supersuit/multiagent_wrappers/black_death.py

@@ -5,24 +5,43 @@
 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)
 
         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
         }

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()