Drawing agent. Code to draw MiniGrid using Matplotlib

gym_minigrid/minigrid.py

@@ -535,7 +535,14 @@ class Grid:
 
         # TODO: overlay agent on top
         if agent_dir is not None:
-            pass
+            tri_fn = point_in_triangle(
+                (0.12, 0.19),
+                (0.87, 0.50),
+                (0.12, 0.81),
+            )
+
+            tri_fn = rotate_fn(tri_fn, cx=0.5, cy=0.5, theta=0.5*math.pi*agent_dir)
+            fill_coords(img, tri_fn, (255, 0, 0))
 
         # TODO: highlighting
         if highlight:
@@ -546,7 +553,12 @@ class Grid:
 
         return img
 
-    def render(self, tile_size):
+    def render(
+        self,
+        tile_size,
+        agent_pos=None,
+        agent_dir=None
+    ):
         """
         Render this grid at a given scale
         :param r: target renderer object
@@ -566,7 +578,7 @@ class Grid:
 
                 tile_img = Grid.render_tile(
                     cell,
-                    agent_dir=None,
+                    agent_dir=agent_dir if agent_pos == (i, j) else None,
                     highlight=False,
                     tile_size=tile_size
                 )
@@ -1302,27 +1314,11 @@ class MiniGridEnv(gym.Env):
         """
 
         # Render the whole grid
-        img = self.grid.render(tile_size)
-
-        """
-        # Draw the agent
-        ratio = tile_size / TILE_PIXELS
-        r.push()
-        r.scale(ratio, ratio)
-        r.translate(
-            TILE_PIXELS * (self.agent_pos[0] + 0.5),
-            TILE_PIXELS * (self.agent_pos[1] + 0.5)
+        img = self.grid.render(
+            tile_size,
+            self.agent_pos,
+            self.agent_dir
         )
-        r.rotate(self.agent_dir * 90)
-        r.setLineColor(255, 0, 0)
-        r.setColor(255, 0, 0)
-        r.drawPolygon([
-            (-12, 10),
-            ( 12,  0),
-            (-12, -10)
-        ])
-        r.pop()
-        """
 
         """
         # Compute which cells are visible to the agent

gym_minigrid/rendering.py

@@ -1,3 +1,4 @@
+import math
 import numpy as np
 
 # TODO: anti-aliased version, fill_coords_aa?
@@ -15,6 +16,18 @@ def fill_coords(img, fn, color):
 
     return img
 
+def rotate_fn(fin, cx, cy, theta):
+    def fout(x, y):
+        x = x - cx
+        y = y - cy
+
+        x2 = cx + x * math.cos(-theta) - y * math.sin(-theta)
+        y2 = cy + y * math.cos(-theta) + x * math.sin(-theta)
+
+        return fin(x2, y2)
+
+    return fout
+
 def point_in_circle(cx, cy, r):
     def fn(x, y):
         return (x-cx)*(x-cx) + (y-cy)*(y-cy) <= r * r

manual_control_matplotlib.py

@@ -0,0 +1,51 @@
+import time
+import matplotlib.pyplot as plt
+import numpy as np
+import gym
+import gym_minigrid
+
+def key_handler(event):
+    print('pressed', event.key)
+
+    if event.key == 'escape':
+        plt.close()
+        return
+
+    if event.key == 'left':
+        env.step(env.actions.left)
+        img = env.render('rgb_array')
+
+        #img = np.zeros(shape=(256,256,3), dtype=np.uint8)
+        imshow_obj.set_data(img)
+        fig.canvas.draw()
+        #plt.show()
+
+        return
+
+env = gym.make('MiniGrid-Empty-8x8-v0')
+
+#env.step(env.actions.left)
+
+
+t0 = time.time()
+
+for i in range(1000):
+    img = env.render('rgb_array')
+
+t1 = time.time()
+dt = int(1000 * (t1-t0))
+
+print(dt)
+
+print(img.shape)
+
+fig, ax = plt.subplots()
+fig.canvas.mpl_connect('key_press_event', key_handler)
+
+#plt.figure(num='gym-minigrid')
+imshow_obj = ax.imshow(img)
+
+
+
+
+plt.show()