diff --git a/gymnasium/utils/env_match.py b/gymnasium/utils/env_match.py index 44182e5fa..fc9f1ecd8 100644 --- a/gymnasium/utils/env_match.py +++ b/gymnasium/utils/env_match.py @@ -27,11 +27,20 @@ def check_environments_match( skip_terminal: If `True` it does not check for equivalence of the observation. skip_truncated: If `True` it does not check for equivalence of the observation. skip_info: If `True` it does not check for equivalence of the observation. - info_comparison: If "equivalence" then checks if the `info`s are identical, - if "superset" checks if `info_b` is a (non-strict) superset of `info_a` - if "skip" no checks are made at the `info`. + info_comparison: The options are + If "equivalence" then checks if the `info`s are identical, + If "superset" checks if `info_b` is a (non-strict) superset of `info_a` + If "keys-equivalence" checks if the `info`s keys are identical (while ignoring the values). + If "keys-superset" checks if the `info_b`s keys are a superset of `info_a`'s keys. + If "skip" no checks are made at the `info`. """ - assert info_comparison in ["equivalence", "superset", "skip"] + assert info_comparison in [ + "equivalence", + "superset", + "skip", + "keys-equivalence", + "keys-superset", + ] assert env_a.action_space == env_b.action_space assert skip_obs or env_b.observation_space == env_b.observation_space @@ -50,6 +59,10 @@ def check_environments_match( assert data_equivalence( info_a[key], info_b[key] ), "resetting info is not a superset" + elif info_comparison == "keys-equivalance": + assert info_a.keys() == info_b.keys(), "resetting info keys are not equivalent" + elif info_comparison == "keys-superset": + assert info_b.keys() >= info_a.keys(), "resetting info keys are not a superset" for _ in range(num_steps): action = env_a.action_space.sample() @@ -74,6 +87,14 @@ def check_environments_match( assert data_equivalence( info_a[key], info_b[key] ), "stepping info is not a superset" + elif info_comparison == "keys-equivalance": + assert ( + info_a.keys() == info_b.keys() + ), "stepping info keys are not equivalent" + elif info_comparison == "keys-superset": + assert ( + info_b.keys() >= info_a.keys() + ), "stepping info keys are not a superset" if terminal_a or truncated_a or terminal_b or truncated_b: obs_a, info_a = env_a.reset(seed=seed) @@ -90,3 +111,11 @@ def check_environments_match( assert data_equivalence( info_a[key], info_b[key] ), "resetting info is not a superset" + elif info_comparison == "keys-equivalance": + assert ( + info_a.keys() == info_b.keys() + ), "resetting info keys are not equivalent" + elif info_comparison == "keys-superset": + assert ( + info_b.keys() >= info_a.keys() + ), "resetting info keys are not a superset"