Browse Source

Implemented box object type

Maxime Chevalier-Boisvert 7 years ago
parent
commit
63d3d94699
3 changed files with 65 additions and 23 deletions
  1. 3 1
      gym_minigrid/envs/fourroomqa.py
  2. 57 22
      gym_minigrid/minigrid.py
  3. 5 0
      gym_minigrid/rendering.py

+ 3 - 1
gym_minigrid/envs/fourroomqa.py

@@ -127,7 +127,7 @@ class FourRoomQAEnv(MiniGridEnv):
         self.startPos = self._randPos(room)
 
         # Possible object types and colors
-        types = ['key', 'ball']
+        types = ['key', 'ball', 'box']
         colors = list(COLORS.keys())
 
         # Place a number of random objects
@@ -140,6 +140,8 @@ class FourRoomQAEnv(MiniGridEnv):
                 obj = Key(objColor)
             elif objType == 'ball':
                 obj = Ball(objColor)
+            elif objType == 'box':
+                obj = Box(objColor)
 
             # Pick a random position that doesn't overlap with anything
             while True:

+ 57 - 22
gym_minigrid/minigrid.py

@@ -42,9 +42,10 @@ OBJECT_TO_IDX = {
     'wall'          : 1,
     'door'          : 2,
     'locked_door'   : 3,
-    'ball'          : 4,
-    'key'           : 5,
-    'goal'          : 6
+    'key'           : 4,
+    'ball'          : 5,
+    'box'           : 6,
+    'goal'          : 7
 }
 
 IDX_TO_OBJECT = dict(zip(OBJECT_TO_IDX.values(), OBJECT_TO_IDX.keys()))
@@ -73,7 +74,7 @@ class WorldObj:
         """Can this contain another object?"""
         return False
 
-    def toggle(self, env):
+    def toggle(self, env, pos):
         """Method to trigger/toggle an action this object performs"""
         return False
 
@@ -144,7 +145,7 @@ class Door(WorldObj):
         ])
         r.drawCircle(CELL_PIXELS * 0.75, CELL_PIXELS * 0.5, 2)
 
-    def toggle(self, env):
+    def toggle(self, env, pos):
         if not self.isOpen:
             self.isOpen = True
             return True
@@ -192,7 +193,7 @@ class LockedDoor(WorldObj):
             CELL_PIXELS * 0.60
         )
 
-    def toggle(self, env):
+    def toggle(self, env, pos):
         # If the player has the right key to open the door
         if isinstance(env.carrying, Key) and env.carrying.color == self.color:
             self.isOpen = True
@@ -205,17 +206,6 @@ class LockedDoor(WorldObj):
         """The agent can only walk over this cell when the door is open"""
         return self.isOpen
 
-class Ball(WorldObj):
-    def __init__(self, color='blue'):
-        super(Ball, self).__init__('ball', color)
-
-    def canPickup(self):
-        return True
-
-    def render(self, r):
-        self._setColor(r)
-        r.drawCircle(CELL_PIXELS * 0.5, CELL_PIXELS * 0.5, 10)
-
 class Key(WorldObj):
     def __init__(self, color='blue'):
         super(Key, self).__init__('key', color)
@@ -253,6 +243,49 @@ class Key(WorldObj):
         r.setColor(0, 0, 0)
         r.drawCircle(18, 9, 2)
 
+class Ball(WorldObj):
+    def __init__(self, color='blue'):
+        super(Ball, self).__init__('ball', color)
+
+    def canPickup(self):
+        return True
+
+    def render(self, r):
+        self._setColor(r)
+        r.drawCircle(CELL_PIXELS * 0.5, CELL_PIXELS * 0.5, 10)
+
+class Box(WorldObj):
+    def __init__(self, color, contains=None):
+        super(Box, self).__init__('box', color)
+        self.contains = contains
+
+    def render(self, r):
+        c = COLORS[self.color]
+        r.setLineColor(c[0], c[1], c[2])
+        r.setColor(0, 0, 0)
+        r.setLineWidth(2)
+
+        r.drawPolygon([
+            (4            , CELL_PIXELS-4),
+            (CELL_PIXELS-4, CELL_PIXELS-4),
+            (CELL_PIXELS-4,             4),
+            (4            ,             4)
+        ])
+
+        r.drawLine(
+            4,
+            CELL_PIXELS / 2,
+            CELL_PIXELS - 4,
+            CELL_PIXELS / 2
+        )
+
+        r.setLineWidth(1)
+
+    def toggle(self, env, pos):
+        # Replace the box by its contents
+        env.grid.set(*pos, self.contains)
+        return True
+
 class Grid:
     """
     Represent a grid and operations on it
@@ -643,12 +676,14 @@ class MiniGridEnv(gym.Env):
         # Pick up or trigger/activate an item
         elif action == self.actions.toggle:
             u, v = self.getDirVec()
-            cell = self.grid.get(self.agentPos[0] + u, self.agentPos[1] + v)
-            if cell and cell.canPickup() and self.carrying is None:
-                self.carrying = cell
-                self.grid.set(self.agentPos[0] + u, self.agentPos[1] + v, None)
+            objPos = (self.agentPos[0] + u, self.agentPos[1] + v)
+            cell = self.grid.get(*objPos)
+            if cell and cell.canPickup():
+                if self.carrying is None:
+                    self.carrying = cell
+                    self.grid.set(*objPos, None)
             elif cell:
-                cell.toggle(self)
+                cell.toggle(self, objPos)
 
         else:
             assert False, "unknown action"

+ 5 - 0
gym_minigrid/rendering.py

@@ -149,6 +149,11 @@ class Renderer:
     def setColor(self, r, g, b, a=255):
         self.painter.setBrush(QColor(r, g, b, a))
 
+    def setLineWidth(self, width):
+        pen = self.painter.pen()
+        pen.setWidthF(width)
+        self.painter.setPen(pen)
+
     def drawLine(self, x0, y0, x1, y1):
         self.painter.drawLine(x0, y0, x1, y1)