From d0196aa3cf365093158fa5b1b850f3bd6f7ee8e9 Mon Sep 17 00:00:00 2001 From: Elliot Tower Date: Wed, 27 Sep 2023 16:41:24 -0400 Subject: [PATCH] Allow parallel envs to have other optional keys in info/obs dicts (e.g., "common") (#1110) --- pettingzoo/test/parallel_test.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pettingzoo/test/parallel_test.py b/pettingzoo/test/parallel_test.py index 56931c198..dd723b328 100644 --- a/pettingzoo/test/parallel_test.py +++ b/pettingzoo/test/parallel_test.py @@ -49,8 +49,9 @@ def parallel_api_test(par_env: ParallelEnv, num_cycles=1000): assert isinstance(obs, dict) assert isinstance(infos, dict) - assert set(obs.keys()) == (set(par_env.agents)) - assert set(infos.keys()) == (set(par_env.agents)) + # Note: obs and info dicts must contain all AgentIDs, but can also have other additional keys (e.g., "common") + assert set(par_env.agents).issubset(set(obs.keys())) + assert set(par_env.agents).issubset(set(infos.keys())) terminated = {agent: False for agent in par_env.agents} truncated = {agent: False for agent in par_env.agents} live_agents = set(par_env.agents[:])