Browse Source

Merge pull request #193 from zeionara/master

Added support for setting a custom max_steps value in doorkey environment
Mark Towers 2 years ago
parent
commit
79f6591174
1 changed files with 8 additions and 8 deletions
  1. 8 8
      gym_minigrid/envs/doorkey.py

+ 8 - 8
gym_minigrid/envs/doorkey.py

@@ -6,10 +6,10 @@ class DoorKeyEnv(MiniGridEnv):
     Environment with a door and key, sparse reward
     """
 
-    def __init__(self, size=8):
+    def __init__(self, size=8, max_steps=None):
         super().__init__(
             grid_size=size,
-            max_steps=10*size*size
+            max_steps=10*size*size if max_steps is None else max_steps
         )
 
     def _gen_grid(self, width, height):
@@ -44,16 +44,16 @@ class DoorKeyEnv(MiniGridEnv):
         self.mission = "use the key to open the door and then get to the goal"
 
 class DoorKeyEnv5x5(DoorKeyEnv):
-    def __init__(self):
-        super().__init__(size=5)
+    def __init__(self, max_steps=None):
+        super().__init__(size=5, max_steps=max_steps)
 
 class DoorKeyEnv6x6(DoorKeyEnv):
-    def __init__(self):
-        super().__init__(size=6)
+    def __init__(self, max_steps=None):
+        super().__init__(size=6, max_steps=max_steps)
 
 class DoorKeyEnv16x16(DoorKeyEnv):
-    def __init__(self):
-        super().__init__(size=16)
+    def __init__(self, max_steps=None):
+        super().__init__(size=16, max_steps=max_steps)
 
 register(
     id='MiniGrid-DoorKey-5x5-v0',