Added "in" operator for Grid objects

gym_minigrid/minigrid.py

@@ -303,6 +303,19 @@ class Grid:
 
         self.grid = [None] * width * height
 
+    def __contains__(self, key):
+        if isinstance(key, WorldObj):
+            for e in self.grid:
+                if e is key:
+                    return True
+        elif isinstance(key, tuple):
+            for e in self.grid:
+                if e is None:
+                    continue
+                if (e.color, e.type) == key:
+                    return True
+        return False
+
     def copy(self):
         from copy import deepcopy
         return deepcopy(self)

run_tests.py

@@ -64,23 +64,19 @@ for envName in sorted(envSet):
 
 ##############################################################################
 
-# Test the env.agentSees() function
 env = gym.make('MiniGrid-Empty-6x6-v0')
 goalPos = (env.grid.width - 2, env.grid.height - 2)
 
-def goalInObs(obs):
-    grid = Grid.decode(obs)
-    for j in range(0, grid.height):
-        for i in range(0, grid.width):
-            cell = grid.get(i, j)
-            if cell and cell.color == 'green':
-                return True
-    return False
+# Test the "in" operator on grid objects
+assert ('green', 'goal') in env.grid
+assert ('blue', 'key') not in env.grid
 
+# Test the env.agentSees() function
 env.reset()
 for i in range(0, 200):
     action = random.randint(0, env.action_space.n - 1)
     obs, reward, done, info = env.step(action)
-    assert env.agentSees(*goalPos) == goalInObs(obs)
+    goalVisible = ('green', 'goal') in Grid.decode(obs)
+    assert env.agentSees(*goalPos) == goalVisible
     if done:
         env.reset()