Split out simple_envs.py

gym_minigrid/__init__.py

@@ -1,3 +1,2 @@
-# Import these environments to make sure they register themselves
-import gym_minigrid.envs.simple_envs
-import gym_minigrid.envs.fourroomqa
+# Import the envs module so that envs register themselves
+import gym_minigrid.envs

gym_minigrid/envs/__init__.py

@@ -1,4 +1,5 @@
-from gym_minigrid.envs.simple_envs import *
+from gym_minigrid.envs.empty import *
+from gym_minigrid.envs.doorkey import *
 from gym_minigrid.envs.multiroom import *
 from gym_minigrid.envs.fetch import *
 from gym_minigrid.envs.fourroomqa import *

gym_minigrid/envs/simple_envs.py

@@ -1,37 +1,6 @@
 from gym_minigrid.minigrid import *
 from gym_minigrid.register import register
 
-class EmptyEnv(MiniGridEnv):
-    """
-    Empty grid environment, no obstacles, sparse reward
-    """
-
-    def __init__(self, size=8):
-        super().__init__(gridSize=size, maxSteps=3 * size)
-
-class EmptyEnv6x6(EmptyEnv):
-    def __init__(self):
-        super().__init__(size=6)
-
-class EmptyEnv16x16(EmptyEnv):
-    def __init__(self):
-        super().__init__(size=16)
-
-register(
-    id='MiniGrid-Empty-6x6-v0',
-    entry_point='gym_minigrid.envs:EmptyEnv6x6'
-)
-
-register(
-    id='MiniGrid-Empty-8x8-v0',
-    entry_point='gym_minigrid.envs:EmptyEnv'
-)
-
-register(
-    id='MiniGrid-Empty-16x16-v0',
-    entry_point='gym_minigrid.envs:EmptyEnv16x16'
-)
-
 class DoorKeyEnv(MiniGridEnv):
     """
     Environment with a door and key, sparse reward

gym_minigrid/envs/empty.py

@@ -0,0 +1,33 @@
+from gym_minigrid.minigrid import *
+from gym_minigrid.register import register
+
+class EmptyEnv(MiniGridEnv):
+    """
+    Empty grid environment, no obstacles, sparse reward
+    """
+
+    def __init__(self, size=8):
+        super().__init__(gridSize=size, maxSteps=3 * size)
+
+class EmptyEnv6x6(EmptyEnv):
+    def __init__(self):
+        super().__init__(size=6)
+
+class EmptyEnv16x16(EmptyEnv):
+    def __init__(self):
+        super().__init__(size=16)
+
+register(
+    id='MiniGrid-Empty-6x6-v0',
+    entry_point='gym_minigrid.envs:EmptyEnv6x6'
+)
+
+register(
+    id='MiniGrid-Empty-8x8-v0',
+    entry_point='gym_minigrid.envs:EmptyEnv'
+)
+
+register(
+    id='MiniGrid-Empty-16x16-v0',
+    entry_point='gym_minigrid.envs:EmptyEnv16x16'
+)

run_tests.py

@@ -6,6 +6,9 @@ import numpy as np
 from gym_minigrid.register import envSet
 from gym_minigrid.minigrid import Grid
 
+# Make sure directly importing a specific environment works
+from gym_minigrid.envs import DoorKeyEnv
+
 print('%d environments registered' % len(envSet))
 
 for envName in sorted(envSet):