add check_environments_match info_comparison options (Farama-Foun…

…dation#584)

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"