浏览代码

Merge pull request #193 from zeionara/master

Added support for setting a custom max_steps value in doorkey environment
Mark Towers 2 年之前
父节点
当前提交
79f6591174
共有 1 个文件被更改,包括 8 次插入8 次删除
  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
     Environment with a door and key, sparse reward
     """
     """
 
 
-    def __init__(self, size=8):
+    def __init__(self, size=8, max_steps=None):
         super().__init__(
         super().__init__(
             grid_size=size,
             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):
     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"
         self.mission = "use the key to open the door and then get to the goal"
 
 
 class DoorKeyEnv5x5(DoorKeyEnv):
 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):
 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):
 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(
 register(
     id='MiniGrid-DoorKey-5x5-v0',
     id='MiniGrid-DoorKey-5x5-v0',