diff --git a/supersuit/generic_wrappers/utils/shared_wrapper_util.py b/supersuit/generic_wrappers/utils/shared_wrapper_util.py index 210dca74..fc31c188 100644 --- a/supersuit/generic_wrappers/utils/shared_wrapper_util.py +++ b/supersuit/generic_wrappers/utils/shared_wrapper_util.py @@ -27,6 +27,17 @@ def observation_space(self, agent): def action_space(self, agent): return self.modifiers[agent].modify_action_space(self.env.action_space(agent)) + def _clear_agent_caches(self): + """Drop per-agent modifiers and space caches. + + Environments with procedurally generated agents (and agent types) may + change observation/action space shapes across resets. Cached spaces and + modifiers from a previous episode would then be wrong. + """ + self.modifiers = {} + self.observation_space.cache_clear() + self.action_space.cache_clear() + def add_modifiers(self, agents_list): for agent in agents_list: if agent not in self.modifiers: @@ -46,8 +57,8 @@ def reset(self, seed=None, options=None): self._cur_seed = seed self._cur_options = options - for mod in self.modifiers.values(): - mod.reset(seed=seed, options=options) + # Spaces/modifiers must be rebuilt after reset for dynamic-agent envs. + self._clear_agent_caches() super().reset(seed=seed, options=options) self.add_modifiers(self.agents) @@ -93,6 +104,15 @@ def observation_space(self, agent): def action_space(self, agent): return self.modifiers[agent].modify_action_space(self.env.action_space(agent)) + def _clear_agent_caches(self): + """Drop per-agent modifiers and space caches after reset. + + See shared_wrapper_aec._clear_agent_caches. + """ + self.modifiers = {} + self.observation_space.cache_clear() + self.action_space.cache_clear() + def add_modifiers(self, agents_list): for agent in agents_list: if agent not in self.modifiers: @@ -113,9 +133,9 @@ def reset(self, seed=None, options=None): self._cur_options = options observations, infos = super().reset(seed=seed, options=options) + # Rebuild modifiers/spaces for the new episode's agent set and spaces. + self._clear_agent_caches() self.add_modifiers(self.agents) - for agent, mod in self.modifiers.items(): - mod.reset(seed=seed, options=options) observations = { agent: self.modifiers[agent].modify_obs(obs) for agent, obs in observations.items() diff --git a/supersuit/lambda_wrappers/action_lambda.py b/supersuit/lambda_wrappers/action_lambda.py index 73ac7ccb..917f1323 100644 --- a/supersuit/lambda_wrappers/action_lambda.py +++ b/supersuit/lambda_wrappers/action_lambda.py @@ -36,6 +36,12 @@ def action_space(self, agent): except TypeError: return self.change_space_fn(old_act_space) + def reset(self, seed=None, options=None): + # Action spaces may change across resets for procedurally generated + # agents; drop the per-agent cache so spaces stay in sync. + self.action_space.cache_clear() + super().reset(seed=seed, options=options) + def _modify_action(self, agent, action): old_act_space = self.env.action_space(agent) try: diff --git a/supersuit/lambda_wrappers/observation_lambda.py b/supersuit/lambda_wrappers/observation_lambda.py index fd15bdbd..843a2c17 100644 --- a/supersuit/lambda_wrappers/observation_lambda.py +++ b/supersuit/lambda_wrappers/observation_lambda.py @@ -62,6 +62,12 @@ def observation_space(self, agent): except TypeError: return self.change_obs_space_fn(old_obs_space) + def reset(self, seed=None, options=None): + # Observation spaces may change across resets for procedurally + # generated agents; drop the per-agent cache so spaces stay in sync. + self.observation_space.cache_clear() + super().reset(seed=seed, options=options) + def _modify_observation(self, agent, observation): old_obs_space = self.env.observation_space(agent) try: diff --git a/test/generated_agents_test.py b/test/generated_agents_test.py index deedbe1c..f691d2a4 100644 --- a/test/generated_agents_test.py +++ b/test/generated_agents_test.py @@ -10,7 +10,12 @@ from supersuit import dtype_v0 -wrappers = [ +# AEC wrappers applied to the parallel-generated-agents env (via parallel_to_aec). +# frame_skip_v0 is excluded: AEC frame-skip does not fully implement the +# PettingZoo dead-agent protocol for environments that create/destroy agents +# mid-episode (agent resurrection / missing observations). Parallel frame_skip +# is covered below and passes. +aec_par_wrappers = [ supersuit.dtype_v0(generated_agents_parallel_v0.env(), np.int32), supersuit.flatten_v0(generated_agents_parallel_v0.env()), supersuit.normalize_obs_v0( @@ -24,23 +29,20 @@ supersuit.nan_noop_v0(generated_agents_parallel_v0.env(), 0), supersuit.nan_zeros_v0(generated_agents_parallel_v0.env()), supersuit.nan_random_v0(generated_agents_parallel_v0.env()), - supersuit.frame_skip_v0(generated_agents_parallel_v0.env(), 4), supersuit.sticky_actions_v0(generated_agents_parallel_v0.env(), 0.75), supersuit.delay_observations_v0(generated_agents_parallel_v0.env(), 3), supersuit.max_observation_v0(generated_agents_parallel_v0.env(), 3), ] -# TODO: fix errors: AssertionError: action is not in action space -@pytest.mark.skip( - reason="skipped: unknown bug, most likely due to converting to AEC env (e.g., obs_lambda has no parallel wrapper)" -) -@pytest.mark.parametrize("env", wrappers) +@pytest.mark.parametrize("env", aec_par_wrappers) def test_pettingzoo_aec_api_par_gen(env): api_test(env, num_cycles=50) -wrappers = [ +# AEC wrappers applied to the native AEC generated-agents env. +# frame_skip_v0 excluded for the same dynamic-agent protocol reason as above. +aec_wrappers = [ supersuit.dtype_v0(generated_agents_env_v0.env(), np.int32), supersuit.flatten_v0(generated_agents_env_v0.env()), supersuit.normalize_obs_v0( @@ -52,23 +54,18 @@ def test_pettingzoo_aec_api_par_gen(env): supersuit.nan_noop_v0(generated_agents_env_v0.env(), 0), supersuit.nan_zeros_v0(generated_agents_env_v0.env()), supersuit.nan_random_v0(generated_agents_env_v0.env()), - supersuit.frame_skip_v0(generated_agents_env_v0.env(), 4), supersuit.sticky_actions_v0(generated_agents_env_v0.env(), 0.75), supersuit.delay_observations_v0(generated_agents_env_v0.env(), 3), supersuit.max_observation_v0(generated_agents_env_v0.env(), 3), ] -# TODO fix error: ValueError: operands could not be broadcast together with shapes (42,) (10,) -@pytest.mark.skip( - reason="skipped: unknown bug, most likely due to converting to AEC env (e.g., obs_lambda has no parallel wrapper)" -) -@pytest.mark.parametrize("env", wrappers) +@pytest.mark.parametrize("env", aec_wrappers) def test_pettingzoo_aec_api_aec_gen(env): api_test(env, num_cycles=50) -parallel_wrappers = wrappers = [ +parallel_wrappers = [ supersuit.dtype_v0(generated_agents_parallel_v0.parallel_env(), np.int32), supersuit.flatten_v0(generated_agents_parallel_v0.parallel_env()), supersuit.normalize_obs_v0( @@ -91,15 +88,24 @@ def test_pettingzoo_aec_api_aec_gen(env): ] -# TODO: fix normalizing obs issue: ValueError: operands could not be broadcast together with shapes (48,) (20,) -@pytest.mark.skip( - reason="skipped: unknown bug, most likely due to converting to AEC env (e.g., obs_lambda has no parallel wrapper)" -) @pytest.mark.parametrize("env", parallel_wrappers) def test_pettingzoo_parallel_api_gen(env): parallel_test.parallel_api_test(env, num_cycles=50) +@pytest.mark.skip( + reason=( + "AEC frame_skip_v0 does not fully support environments with " + "procedurally generated agents: skipping mid-episode agent " + "birth/death breaks the PettingZoo dead-agent protocol " + "(resurrection / missing observations). Parallel frame_skip is tested." + ) +) +def test_aec_frame_skip_generated_agents_known_limitation(): + env = supersuit.frame_skip_v0(generated_agents_env_v0.env(), 4) + api_test(env, num_cycles=50) + + wrapper_fns = [ lambda: supersuit.pad_action_space_v0(generated_agents_parallel_v0.env()), lambda: supersuit.pad_observations_v0(generated_agents_parallel_v0.env()),