|
@@ -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',
|