Fixed key in Dict.keys() always being false (#608)

Author: pseudo-rnd-thoughts

gymnasium/spaces/dict.py

@@ -4,7 +4,7 @@
 import collections.abc
 import typing
 from collections import OrderedDict
-from typing import Any, Sequence
+from typing import Any, KeysView, Sequence
 
 import numpy as np
 
@@ -180,6 +180,10 @@ def __getitem__(self, key: str) -> Space[Any]:
         """Get the space that is associated to `key`."""
         return self.spaces[key]
 
+    def keys(self) -> KeysView:
+        """Returns the keys of the Dict."""
+        return KeysView(self.spaces)
+
     def __setitem__(self, key: str, value: Space[Any]):
         """Set the space that is associated to `key`."""
         assert isinstance(

tests/spaces/test_dict.py

@@ -145,3 +145,14 @@ def test_iterator():
         assert key in DICT_SPACE.spaces
 
     assert {key for key in DICT_SPACE} == DICT_SPACE.spaces.keys()
+
+
+def test_keys_contains():
+    """Test that `Dict.keys()` will correctly assess if the key is in the space."""
+    space = Dict(a=Box(0, 1), b=Box(1, 2))
+
+    for key in space.keys():
+        assert key in space.keys()
+    assert "a" in space.keys()
+
+    assert "c" not in space.keys()