Cleaned up and simplified _genGrid functions

gym_minigrid/envs/doorkey.py

@@ -11,23 +11,20 @@ class DoorKeyEnv(MiniGridEnv):
 
     def _genGrid(self, width, height):
         # Create an empty grid
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
 
-        # Place walls around the edges
-        for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height - 1, Wall())
-        for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(height - 1, j, Wall())
+        # Generate the surrounding walls
+        self.grid.horzWall(0, 0)
+        self.grid.horzWall(0, height-1)
+        self.grid.vertWall(0, 0)
+        self.grid.vertWall(width-1, 0)
 
         # Place a goal in the bottom-right corner
-        grid.set(width - 2, height - 2, Goal())
+        self.grid.set(width - 2, height - 2, Goal())
 
         # Create a vertical splitting wall
         splitIdx = self._randInt(2, width-2)
-        for i in range(0, height):
-            grid.set(splitIdx, i, Wall())
+        self.grid.vertWall(splitIdx, 0)
 
         # Place the agent at a random position and orientation
         self.startPos = self._randPos(
@@ -38,7 +35,7 @@ class DoorKeyEnv(MiniGridEnv):
 
         # Place a door in the wall
         doorIdx = self._randInt(1, width-2)
-        grid.set(splitIdx, doorIdx, LockedDoor('yellow'))
+        self.grid.set(splitIdx, doorIdx, LockedDoor('yellow'))
 
         # Place a yellow key on the left side
         while True:
@@ -48,15 +45,13 @@ class DoorKeyEnv(MiniGridEnv):
             )
             if pos == self.startPos:
                 continue
-            if grid.get(*pos) != None:
+            if self.grid.get(*pos) != None:
                 continue
-            grid.set(*pos, Key('yellow'))
+            self.grid.set(*pos, Key('yellow'))
             break
 
         self.mission = "use the key to open the door and then get to the goal"
 
-        return grid
-
 class DoorKeyEnv5x5(DoorKeyEnv):
     def __init__(self):
         super().__init__(size=5)

gym_minigrid/envs/empty.py

@@ -11,23 +11,19 @@ class EmptyEnv(MiniGridEnv):
 
     def _genGrid(self, width, height):
         # Create an empty grid
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
 
-        # Place walls around the edges
-        for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height - 1, Wall())
-        for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(height - 1, j, Wall())
+        # Generate the surrounding walls
+        self.grid.horzWall(0, 0)
+        self.grid.horzWall(0, height-1)
+        self.grid.vertWall(0, 0)
+        self.grid.vertWall(width-1, 0)
 
-        # Place a goal in the bottom-right corner
-        grid.set(width - 2, height - 2, Goal())
+        # Place a goal square in the bottom-right corner
+        self.grid.set(width - 2, height - 2, Goal())
 
         self.mission = "get to the green goal square"
 
-        return grid
-
 class EmptyEnv6x6(EmptyEnv):
     def __init__(self):
         super().__init__(size=6)

gym_minigrid/envs/fetch.py

@@ -17,45 +17,34 @@ class FetchEnv(MiniGridEnv):
         self.reward_range = (0, 1)
 
     def _genGrid(self, width, height):
-        assert width == height
-        gridSz = width
-
-        # Create a grid surrounded by walls
-        grid = Grid(width, height)
-        for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height-1, Wall())
-        for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(width-1, j, Wall())
+        self.grid = Grid(width, height)
+
+        # Generate the surrounding walls
+        self.grid.horzWall(0, 0)
+        self.grid.horzWall(0, height-1)
+        self.grid.vertWall(0, 0)
+        self.grid.vertWall(width-1, 0)
 
         types = ['key', 'ball']
-        colors = list(COLORS.keys())
 
         objs = []
 
         # For each object to be generated
         while len(objs) < self.numObjs:
             objType = self._randElem(types)
-            objColor = self._randElem(colors)
+            objColor = self._randElem(COLOR_NAMES)
 
             if objType == 'key':
                 obj = Key(objColor)
             elif objType == 'ball':
                 obj = Ball(objColor)
 
-            while True:
-                pos = (
-                    self._randInt(1, gridSz - 1),
-                    self._randInt(1, gridSz - 1)
-                )
-
-                if pos != self.startPos:
-                    grid.set(*pos, obj)
-                    break
-
+            self.placeObj(obj)
             objs.append(obj)
 
+        # Randomize the player start position and orientation
+        self.placeAgent()
+
         # Choose a random object to be picked up
         target = objs[self._randInt(0, len(objs))]
         self.targetType = target.type
@@ -77,8 +66,6 @@ class FetchEnv(MiniGridEnv):
             self.mission = 'you must fetch a %s' % descStr
         assert hasattr(self, 'mission')
 
-        return grid
-
     def step(self, action):
         obs, reward, done, info = MiniGridEnv.step(self, action)
 

gym_minigrid/envs/fourroomqa.py

@@ -59,7 +59,7 @@ class FourRoomQAEnv(MiniGridEnv):
         )
 
     def _genGrid(self, width, height):
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
 
         # Horizontal and vertical split indices
         vSplitIdx = self._randInt(5, width-4)
@@ -99,28 +99,28 @@ class FourRoomQAEnv(MiniGridEnv):
 
             # Horizontal walls
             for i in range(w):
-                grid.set(x + i, y, Wall(room.color))
-                grid.set(x + i, y + h - 1, Wall(room.color))
+                self.grid.set(x + i, y, Wall(room.color))
+                self.grid.set(x + i, y + h - 1, Wall(room.color))
 
             # Vertical walls
             for j in range(h):
-                grid.set(x, y + j, Wall(room.color))
-                grid.set(x + w - 1, y + j, Wall(room.color))
+                self.grid.set(x, y + j, Wall(room.color))
+                self.grid.set(x + w - 1, y + j, Wall(room.color))
 
         # Place wall openings connecting the rooms
         hIdx = self._randInt(1, hSplitIdx-1)
-        grid.set(vSplitIdx, hIdx, None)
-        grid.set(vSplitIdx-1, hIdx, None)
+        self.grid.set(vSplitIdx, hIdx, None)
+        self.grid.set(vSplitIdx-1, hIdx, None)
         hIdx = self._randInt(hSplitIdx+1, height-1)
-        grid.set(vSplitIdx, hIdx, None)
-        grid.set(vSplitIdx-1, hIdx, None)
+        self.grid.set(vSplitIdx, hIdx, None)
+        self.grid.set(vSplitIdx-1, hIdx, None)
 
         vIdx = self._randInt(1, vSplitIdx-1)
-        grid.set(vIdx, hSplitIdx, None)
-        grid.set(vIdx, hSplitIdx-1, None)
+        self.grid.set(vIdx, hSplitIdx, None)
+        self.grid.set(vIdx, hSplitIdx-1, None)
         vIdx = self._randInt(vSplitIdx+1, width-1)
-        grid.set(vIdx, hSplitIdx, None)
-        grid.set(vIdx, hSplitIdx-1, None)
+        self.grid.set(vIdx, hSplitIdx, None)
+        self.grid.set(vIdx, hSplitIdx-1, None)
 
         # Select a random position for the agent to start at
         self.startDir = self._randInt(0, 4)
@@ -150,9 +150,9 @@ class FourRoomQAEnv(MiniGridEnv):
                 pos = self._randPos(room, border=2)
                 if pos == self.startPos:
                     continue
-                if grid.get(*pos) != None:
+                if self.grid.get(*pos) != None:
                     continue
-                grid.set(*pos, obj)
+                self.grid.set(*pos, obj)
                 break
 
             room.objects.append(obj)
@@ -182,11 +182,6 @@ class FourRoomQAEnv(MiniGridEnv):
 
         # TODO: how many X in the Y room question type
 
-        #print(self.mission)
-        #print(self.answer)
-
-        return grid
-
     def step(self, action):
         if isinstance(action, dict):
             answer = action['answer']

gym_minigrid/envs/gotodoor.py

@@ -17,26 +17,17 @@ class GoToDoorEnv(MiniGridEnv):
 
     def _genGrid(self, width, height):
         # Create the grid
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
 
         # Randomly vary the room width and height
         width = self._randInt(5, width+1)
         height = self._randInt(5, height+1)
 
         # Generate the surrounding walls
-        for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height-1, Wall())
-        for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(width-1, j, Wall())
-
-        # Randomize the player start position and orientation
-        self.startPos = self._randPos(
-            1, width-1,
-            1, height-1
-        )
-        self.startDir = self._randInt(0, 4)
+        self.grid.horzWall(0, 0, width)
+        self.grid.horzWall(0, height-1, width)
+        self.grid.vertWall(0, 0, height)
+        self.grid.vertWall(width-1, 0, height)
 
         # Generate the 4 doors at random positions
         doorPos = []
@@ -56,7 +47,7 @@ class GoToDoorEnv(MiniGridEnv):
         # Place the doors in the grid
         for idx, pos in enumerate(doorPos):
             color = doorColors[idx]
-            grid.set(*pos, Door(color))
+            self.grid.set(*pos, Door(color))
 
         # Select a random target door
         doorIdx = self._randInt(0, len(doorPos))
@@ -65,9 +56,6 @@ class GoToDoorEnv(MiniGridEnv):
 
         # Generate the mission string
         self.mission = 'go to the %s door' % self.targetColor
-        #print(self.mission)
-
-        return grid
 
     def step(self, action):
         obs, reward, done, info = MiniGridEnv.step(self, action)

gym_minigrid/envs/gotoobject.py

@@ -17,17 +17,13 @@ class GoToObjectEnv(MiniGridEnv):
         self.reward_range = (0, 1)
 
     def _genGrid(self, width, height):
-        assert width == height
-        gridSz = width
-
-        # Create a grid surrounded by walls
-        grid = Grid(width, height)
-        for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height-1, Wall())
-        for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(width-1, j, Wall())
+        self.grid = Grid(width, height)
+
+        # Generate the surrounding walls
+        self.grid.horzWall(0, 0)
+        self.grid.horzWall(0, height-1)
+        self.grid.vertWall(0, 0)
+        self.grid.vertWall(width-1, 0)
 
         # Types and colors of objects we can generate
         types = ['key', 'ball', 'box']
@@ -51,10 +47,13 @@ class GoToObjectEnv(MiniGridEnv):
             elif objType == 'box':
                 obj = Box(objColor)
 
-            pos = self.placeObj(grid, obj, self.startPos)
+            pos = self.placeObj(obj)
             objs.append((objType, objColor))
             objPos.append(pos)
 
+        # Randomize the player start position and orientation
+        self.placeAgent()
+
         # Choose a random object to be picked up
         objIdx = self._randInt(0, len(objs))
         self.targetType, self.targetColor = objs[objIdx]
@@ -64,8 +63,6 @@ class GoToObjectEnv(MiniGridEnv):
         self.mission = 'go to the %s' % descStr
         #print(self.mission)
 
-        return grid
-
     def step(self, action):
         obs, reward, done, info = MiniGridEnv.step(self, action)
 

gym_minigrid/envs/lockedroom.py

@@ -40,22 +40,22 @@ class LockedRoom(MiniGridEnv):
 
     def _genGrid(self, width, height):
         # Create the grid
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
 
         # Generate the surrounding walls
         for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height-1, Wall())
+            self.grid.set(i, 0, Wall())
+            self.grid.set(i, height-1, Wall())
         for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(width-1, j, Wall())
+            self.grid.set(0, j, Wall())
+            self.grid.set(width-1, j, Wall())
 
         # Hallway walls
         lWallIdx = width // 2 - 2
         rWallIdx = width // 2 + 2
         for j in range(0, height):
-            grid.set(lWallIdx, j, Wall())
-            grid.set(rWallIdx, j, Wall())
+            self.grid.set(lWallIdx, j, Wall())
+            self.grid.set(rWallIdx, j, Wall())
 
         self.rooms = []
 
@@ -63,9 +63,9 @@ class LockedRoom(MiniGridEnv):
         for n in range(0, 3):
             j = n * (height // 3)
             for i in range(0, lWallIdx):
-                grid.set(i, j, Wall())
+                self.grid.set(i, j, Wall())
             for i in range(rWallIdx, width):
-                grid.set(i, j, Wall())
+                self.grid.set(i, j, Wall())
 
             roomW = lWallIdx + 1
             roomH = height // 3 + 1
@@ -84,7 +84,7 @@ class LockedRoom(MiniGridEnv):
         lockedRoom = self._randElem(self.rooms)
         lockedRoom.locked = True
         goalPos = lockedRoom.randPos(self)
-        grid.set(*goalPos, Goal())
+        self.grid.set(*goalPos, Goal())
 
         # Assign the door colors
         colors = set(COLOR_NAMES)
@@ -93,9 +93,9 @@ class LockedRoom(MiniGridEnv):
             colors.remove(color)
             room.color = color
             if room.locked:
-                grid.set(*room.doorPos, LockedDoor(color))
+                self.grid.set(*room.doorPos, LockedDoor(color))
             else:
-                grid.set(*room.doorPos, Door(color))
+                self.grid.set(*room.doorPos, Door(color))
 
         # Select a random room to contain the key
         while True:
@@ -103,7 +103,7 @@ class LockedRoom(MiniGridEnv):
             if keyRoom != lockedRoom:
                 break
         keyPos = keyRoom.randPos(self)
-        grid.set(*keyPos, Key(lockedRoom.color))
+        self.grid.set(*keyPos, Key(lockedRoom.color))
 
         # Randomize the player start position and orientation
         self.startPos = self._randPos(
@@ -119,8 +119,6 @@ class LockedRoom(MiniGridEnv):
             'so you can get to the goal'
         ) % (lockedRoom.color, keyRoom.color, lockedRoom.color)
 
-        return grid
-
     def step(self, action):
         obs, reward, done, info = MiniGridEnv.step(self, action)
         return obs, reward, done, info

gym_minigrid/envs/multiroom.py

@@ -80,7 +80,7 @@ class MultiRoomEnv(MiniGridEnv):
         self.startDir = self._randInt(0, 4)
 
         # Create the grid
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
         wall = Wall()
 
         prevDoorColor = None
@@ -93,13 +93,13 @@ class MultiRoomEnv(MiniGridEnv):
 
             # Draw the top and bottom walls
             for i in range(0, sizeX):
-                grid.set(topX + i, topY, wall)
-                grid.set(topX + i, topY + sizeY - 1, wall)
+                self.grid.set(topX + i, topY, wall)
+                self.grid.set(topX + i, topY + sizeY - 1, wall)
 
             # Draw the left and right walls
             for j in range(0, sizeY):
-                grid.set(topX, topY + j, wall)
-                grid.set(topX + sizeX - 1, topY + j, wall)
+                self.grid.set(topX, topY + j, wall)
+                self.grid.set(topX + sizeX - 1, topY + j, wall)
 
             # If this isn't the first room, place the entry door
             if idx > 0:
@@ -112,7 +112,7 @@ class MultiRoomEnv(MiniGridEnv):
                 doorColor = self._randElem(sorted(doorColors))
 
                 entryDoor = Door(doorColor)
-                grid.set(*room.entryDoorPos, entryDoor)
+                self.grid.set(*room.entryDoorPos, entryDoor)
                 prevDoorColor = doorColor
 
                 prevRoom = roomList[idx-1]
@@ -127,13 +127,11 @@ class MultiRoomEnv(MiniGridEnv):
 
             # Make sure the goal doesn't overlap with the agent
             if self.goalPos != self.startPos:
-                grid.set(*self.goalPos, Goal())
+                self.grid.set(*self.goalPos, Goal())
                 break
 
         self.mission = 'traverse the rooms to get to the goal'
 
-        return grid
-
     def _placeRoom(
         self,
         numLeft,

gym_minigrid/envs/playground_v0.py

@@ -13,22 +13,17 @@ class PlaygroundV0(MiniGridEnv):
 
     def _genGrid(self, width, height):
         # Create the grid
-        grid = Grid(width, height)
+        self.grid = Grid(width, height)
 
         # Generate the surrounding walls
-        grid.horzWall(0, 0)
-        grid.horzWall(0, height-1)
-        grid.vertWall(0, 0)
-        grid.vertWall(width-1, 0)
+        self.grid.horzWall(0, 0)
+        self.grid.horzWall(0, height-1)
+        self.grid.vertWall(0, 0)
+        self.grid.vertWall(width-1, 0)
 
         roomW = width // 3
         roomH = height // 3
 
-        # NOTE: if we want no room to have two doors of the same color,
-        # that adds difficulty. It's a graph coloring problem
-        # We could generate the door positions first, and do the coloring
-        # in a second pass through rejection sampling.
-
         # For each row of rooms
         for j in range(0, 3):
 
@@ -41,20 +36,20 @@ class PlaygroundV0(MiniGridEnv):
 
                 # Bottom wall and door
                 if i+1 < 3:
-                    grid.vertWall(xR, yT, roomH)
+                    self.grid.vertWall(xR, yT, roomH)
                     pos = (xR, self._randInt(yT+1, yB-1))
                     color = self._randElem(COLOR_NAMES)
-                    grid.set(*pos, Door(color))
+                    self.grid.set(*pos, Door(color))
 
                 # Bottom wall and door
                 if j+1 < 3:
-                    grid.horzWall(xL, yB, roomW)
+                    self.grid.horzWall(xL, yB, roomW)
                     pos = (self._randInt(xL+1, xR-1), yB)
                     color = self._randElem(COLOR_NAMES)
-                    grid.set(*pos, Door(color))
+                    self.grid.set(*pos, Door(color))
 
         # Randomize the player start position and orientation
-        self.placeAgent(grid)
+        self.placeAgent()
 
         # Place random objects in the world
         types = ['key', 'ball', 'box']
@@ -67,13 +62,11 @@ class PlaygroundV0(MiniGridEnv):
                 obj = Ball(objColor)
             elif objType == 'box':
                 obj = Box(objColor)
-            self.placeObj(grid, obj, self.startPos)
+            self.placeObj(obj)
 
         # No explicit mission in this environment
         self.mission = ''
 
-        return grid
-
     def step(self, action):
         obs, reward, done, info = MiniGridEnv.step(self, action)
         return obs, reward, done, info

gym_minigrid/envs/putnear.py

@@ -17,14 +17,13 @@ class PutNearEnv(MiniGridEnv):
         self.reward_range = (0, 1)
 
     def _genGrid(self, width, height):
-        # Create a grid surrounded by walls
-        grid = Grid(width, height)
-        for i in range(0, width):
-            grid.set(i, 0, Wall())
-            grid.set(i, height-1, Wall())
-        for j in range(0, height):
-            grid.set(0, j, Wall())
-            grid.set(width-1, j, Wall())
+        self.grid = Grid(width, height)
+
+        # Generate the surrounding walls
+        self.grid.horzWall(0, 0)
+        self.grid.horzWall(0, height-1)
+        self.grid.vertWall(0, 0)
+        self.grid.vertWall(width-1, 0)
 
         # Types and colors of objects we can generate
         types = ['key', 'ball']
@@ -66,7 +65,7 @@ class PutNearEnv(MiniGridEnv):
                     continue
                 if pos == self.startPos:
                     continue
-                grid.set(*pos, obj)
+                self.grid.set(*pos, obj)
                 break
 
             objs.append((objType, objColor))
@@ -92,8 +91,6 @@ class PutNearEnv(MiniGridEnv):
             self.targetType
         )
 
-        return grid
-
     def step(self, action):
         preCarrying = self.carrying
 

gym_minigrid/minigrid.py

@@ -577,7 +577,7 @@ class MiniGridEnv(gym.Env):
         # Generate a new random grid at the start of each episode
         # To keep the same grid for each episode, call env.seed() with
         # the same seed before calling env.reset()
-        self.grid = self._genGrid(self.gridSize, self.gridSize)
+        self._genGrid(self.gridSize, self.gridSize)
 
         # Place the agent in the starting position and direction
         self.agentPos = self.startPos
@@ -606,6 +606,15 @@ class MiniGridEnv(gym.Env):
 
         return self.np_random.randint(low, high)
 
+    def _randElem(self, iterable):
+        """
+        Pick a random element in a list
+        """
+
+        lst = list(iterable)
+        idx = self._randInt(0, len(lst))
+        return lst[idx]
+
     def _randPos(self, xLow, xHigh, yLow, yHigh):
         """
         Generate a random (x,y) position tuple
@@ -616,22 +625,30 @@ class MiniGridEnv(gym.Env):
             self.np_random.randint(yLow, yHigh)
         )
 
-    def placeObj(self, grid, obj, excPos=None):
+    def placeObj(self, obj):
+        """
+        Place an object at an empty position in the grid
+        """
+
         while True:
             pos = (
-                self._randInt(0, grid.width),
-                self._randInt(0, grid.height)
+                self._randInt(0, self.grid.width),
+                self._randInt(0, self.grid.height)
             )
-            if grid.get(*pos) != None:
+            if self.grid.get(*pos) != None:
                 continue
-            if pos == excPos:
+            if pos == self.startPos:
                 continue
             break
-        grid.set(*pos, obj)
+        self.grid.set(*pos, obj)
         return pos
 
-    def placeAgent(self, grid, randDir=True):
-        pos = self.placeObj(grid, None)
+    def placeAgent(self, randDir=True):
+        """
+        Set the agent's starting point at an empty position in the grid
+        """
+
+        pos = self.placeObj(None)
         self.startPos = pos
 
         if randDir:
@@ -639,11 +656,6 @@ class MiniGridEnv(gym.Env):
 
         return pos
 
-    def _randElem(self, iterable):
-        lst = list(iterable)
-        idx = self._randInt(0, len(lst))
-        return lst[idx]
-
     def getStepsRemaining(self):
         return self.maxSteps - self.stepCount