minigrid.py 36 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328
  1. import math
  2. import gym
  3. from enum import IntEnum
  4. import numpy as np
  5. from gym import error, spaces, utils
  6. from gym.utils import seeding
  7. from .rendering import *
  8. # Size in pixels of a tile in the full-scale human view
  9. TILE_PIXELS = 32
  10. # Map of color names to RGB values
  11. COLORS = {
  12. 'red' : np.array([255, 0, 0]),
  13. 'green' : np.array([0, 255, 0]),
  14. 'blue' : np.array([0, 0, 255]),
  15. 'purple': np.array([112, 39, 195]),
  16. 'yellow': np.array([255, 255, 0]),
  17. 'grey' : np.array([100, 100, 100])
  18. }
  19. COLOR_NAMES = sorted(list(COLORS.keys()))
  20. # Used to map colors to integers
  21. COLOR_TO_IDX = {
  22. 'red' : 0,
  23. 'green' : 1,
  24. 'blue' : 2,
  25. 'purple': 3,
  26. 'yellow': 4,
  27. 'grey' : 5
  28. }
  29. IDX_TO_COLOR = dict(zip(COLOR_TO_IDX.values(), COLOR_TO_IDX.keys()))
  30. # Map of object type to integers
  31. OBJECT_TO_IDX = {
  32. 'unseen' : 0,
  33. 'empty' : 1,
  34. 'wall' : 2,
  35. 'floor' : 3,
  36. 'door' : 4,
  37. 'key' : 5,
  38. 'ball' : 6,
  39. 'box' : 7,
  40. 'goal' : 8,
  41. 'lava' : 9,
  42. 'agent' : 10,
  43. }
  44. IDX_TO_OBJECT = dict(zip(OBJECT_TO_IDX.values(), OBJECT_TO_IDX.keys()))
  45. # Map of state names to integers
  46. STATE_TO_IDX = {
  47. 'open' : 0,
  48. 'closed': 1,
  49. 'locked': 2,
  50. }
  51. # Map of agent direction indices to vectors
  52. DIR_TO_VEC = [
  53. # Pointing right (positive X)
  54. np.array((1, 0)),
  55. # Down (positive Y)
  56. np.array((0, 1)),
  57. # Pointing left (negative X)
  58. np.array((-1, 0)),
  59. # Up (negative Y)
  60. np.array((0, -1)),
  61. ]
  62. class WorldObj:
  63. """
  64. Base class for grid world objects
  65. """
  66. def __init__(self, type, color):
  67. assert type in OBJECT_TO_IDX, type
  68. assert color in COLOR_TO_IDX, color
  69. self.type = type
  70. self.color = color
  71. self.contains = None
  72. # Initial position of the object
  73. self.init_pos = None
  74. # Current position of the object
  75. self.cur_pos = None
  76. def can_overlap(self):
  77. """Can the agent overlap with this?"""
  78. return False
  79. def can_pickup(self):
  80. """Can the agent pick this up?"""
  81. return False
  82. def can_contain(self):
  83. """Can this contain another object?"""
  84. return False
  85. def see_behind(self):
  86. """Can the agent see behind this object?"""
  87. return True
  88. def toggle(self, env, pos):
  89. """Method to trigger/toggle an action this object performs"""
  90. return False
  91. def encode(self):
  92. """Encode the a description of this object as a 3-tuple of integers"""
  93. return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], 0)
  94. @staticmethod
  95. def decode(type_idx, color_idx, state):
  96. """Create an object from a 3-tuple state description"""
  97. obj_type = IDX_TO_OBJECT[type_idx]
  98. color = IDX_TO_COLOR[color_idx]
  99. if obj_type == 'empty' or obj_type == 'unseen':
  100. return None
  101. # State, 0: open, 1: closed, 2: locked
  102. is_open = state == 0
  103. is_locked = state == 2
  104. if obj_type == 'wall':
  105. v = Wall(color)
  106. elif obj_type == 'floor':
  107. v = Floor(color)
  108. elif obj_type == 'ball':
  109. v = Ball(color)
  110. elif obj_type == 'key':
  111. v = Key(color)
  112. elif obj_type == 'box':
  113. v = Box(color)
  114. elif obj_type == 'door':
  115. v = Door(color, is_open, is_locked)
  116. elif obj_type == 'goal':
  117. v = Goal()
  118. elif obj_type == 'lava':
  119. v = Lava()
  120. else:
  121. assert False, "unknown object type in decode '%s'" % objType
  122. return v
  123. def render(self, r):
  124. """Draw this object with the given renderer"""
  125. raise NotImplementedError
  126. class Goal(WorldObj):
  127. def __init__(self):
  128. super().__init__('goal', 'green')
  129. def can_overlap(self):
  130. return True
  131. def render(self, img):
  132. fill_coords(img, point_in_rect(0, 1, 0, 1), COLORS[self.color])
  133. class Floor(WorldObj):
  134. """
  135. Colored floor tile the agent can walk over
  136. """
  137. def __init__(self, color='blue'):
  138. super().__init__('floor', color)
  139. def can_overlap(self):
  140. return True
  141. def render(self, r):
  142. # Give the floor a pale color
  143. c = COLORS[self.color]
  144. r.setLineColor(100, 100, 100, 0)
  145. r.setColor(*c/2)
  146. r.drawPolygon([
  147. (1 , TILE_PIXELS),
  148. (TILE_PIXELS, TILE_PIXELS),
  149. (TILE_PIXELS, 1),
  150. (1 , 1)
  151. ])
  152. class Lava(WorldObj):
  153. def __init__(self):
  154. super().__init__('lava', 'red')
  155. def can_overlap(self):
  156. return True
  157. def render(self, r):
  158. orange = 255, 128, 0
  159. r.setLineColor(*orange)
  160. r.setColor(*orange)
  161. r.drawPolygon([
  162. (0 , TILE_PIXELS),
  163. (TILE_PIXELS, TILE_PIXELS),
  164. (TILE_PIXELS, 0),
  165. (0 , 0)
  166. ])
  167. # drawing the waves
  168. r.setLineColor(0, 0, 0)
  169. r.drawPolyline([
  170. (.1 * TILE_PIXELS, .3 * TILE_PIXELS),
  171. (.3 * TILE_PIXELS, .4 * TILE_PIXELS),
  172. (.5 * TILE_PIXELS, .3 * TILE_PIXELS),
  173. (.7 * TILE_PIXELS, .4 * TILE_PIXELS),
  174. (.9 * TILE_PIXELS, .3 * TILE_PIXELS),
  175. ])
  176. r.drawPolyline([
  177. (.1 * TILE_PIXELS, .5 * TILE_PIXELS),
  178. (.3 * TILE_PIXELS, .6 * TILE_PIXELS),
  179. (.5 * TILE_PIXELS, .5 * TILE_PIXELS),
  180. (.7 * TILE_PIXELS, .6 * TILE_PIXELS),
  181. (.9 * TILE_PIXELS, .5 * TILE_PIXELS),
  182. ])
  183. r.drawPolyline([
  184. (.1 * TILE_PIXELS, .7 * TILE_PIXELS),
  185. (.3 * TILE_PIXELS, .8 * TILE_PIXELS),
  186. (.5 * TILE_PIXELS, .7 * TILE_PIXELS),
  187. (.7 * TILE_PIXELS, .8 * TILE_PIXELS),
  188. (.9 * TILE_PIXELS, .7 * TILE_PIXELS),
  189. ])
  190. class Wall(WorldObj):
  191. def __init__(self, color='grey'):
  192. super().__init__('wall', color)
  193. def see_behind(self):
  194. return False
  195. def render(self, img):
  196. fill_coords(img, point_in_rect(0, 1, 0, 1), COLORS[self.color])
  197. class Door(WorldObj):
  198. def __init__(self, color, is_open=False, is_locked=False):
  199. super().__init__('door', color)
  200. self.is_open = is_open
  201. self.is_locked = is_locked
  202. def can_overlap(self):
  203. """The agent can only walk over this cell when the door is open"""
  204. return self.is_open
  205. def see_behind(self):
  206. return self.is_open
  207. def toggle(self, env, pos):
  208. # If the player has the right key to open the door
  209. if self.is_locked:
  210. if isinstance(env.carrying, Key) and env.carrying.color == self.color:
  211. self.is_locked = False
  212. self.is_open = True
  213. return True
  214. return False
  215. self.is_open = not self.is_open
  216. return True
  217. def encode(self):
  218. """Encode the a description of this object as a 3-tuple of integers"""
  219. # State, 0: open, 1: closed, 2: locked
  220. if self.is_open:
  221. state = 0
  222. elif self.is_locked:
  223. state = 2
  224. elif not self.is_open:
  225. state = 1
  226. return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], state)
  227. def render(self, img):
  228. c = COLORS[self.color]
  229. if self.is_open:
  230. fill_coords(img, point_in_rect(0.88, 1.00, 0.00, 1.00), c)
  231. fill_coords(img, point_in_rect(0.92, 0.96, 0.04, 0.96), (0,0,0))
  232. return
  233. # Door frame and door
  234. if self.is_locked:
  235. fill_coords(img, point_in_rect(0.00, 1.00, 0.00, 1.00), c)
  236. fill_coords(img, point_in_rect(0.06, 0.94, 0.06, 0.94), 0.45 * np.array(c))
  237. # Draw key slot
  238. fill_coords(img, point_in_rect(0.52, 0.75, 0.50, 0.56), c)
  239. else:
  240. fill_coords(img, point_in_rect(0.00, 1.00, 0.00, 1.00), c)
  241. fill_coords(img, point_in_rect(0.04, 0.96, 0.04, 0.96), (0,0,0))
  242. fill_coords(img, point_in_rect(0.08, 0.92, 0.08, 0.92), c)
  243. fill_coords(img, point_in_rect(0.12, 0.88, 0.12, 0.88), (0,0,0))
  244. # Draw door handle
  245. fill_coords(img, point_in_circle(cx=0.75, cy=0.50, r=0.08), c)
  246. class Key(WorldObj):
  247. def __init__(self, color='blue'):
  248. super(Key, self).__init__('key', color)
  249. def can_pickup(self):
  250. return True
  251. def render(self, img):
  252. c = COLORS[self.color]
  253. # Vertical quad
  254. fill_coords(img, point_in_rect(0.50, 0.63, 0.31, 0.89), c)
  255. # Teeth
  256. fill_coords(img, point_in_rect(0.38, 0.50, 0.59, 0.66), c)
  257. fill_coords(img, point_in_rect(0.38, 0.50, 0.81, 0.88), c)
  258. # Ring
  259. fill_coords(img, point_in_circle(cx=0.56, cy=0.28, r=0.190), c)
  260. fill_coords(img, point_in_circle(cx=0.56, cy=0.28, r=0.064), (0,0,0))
  261. class Ball(WorldObj):
  262. def __init__(self, color='blue'):
  263. super(Ball, self).__init__('ball', color)
  264. def can_pickup(self):
  265. return True
  266. def render(self, img):
  267. fill_coords(img, point_in_circle(0.5, 0.5, 0.31), COLORS[self.color])
  268. class Box(WorldObj):
  269. def __init__(self, color, contains=None):
  270. super(Box, self).__init__('box', color)
  271. self.contains = contains
  272. def can_pickup(self):
  273. return True
  274. def render(self, r):
  275. c = COLORS[self.color]
  276. r.setLineColor(c[0], c[1], c[2])
  277. r.setColor(0, 0, 0)
  278. r.setLineWidth(2)
  279. r.drawPolygon([
  280. (4 , TILE_PIXELS-4),
  281. (TILE_PIXELS-4, TILE_PIXELS-4),
  282. (TILE_PIXELS-4, 4),
  283. (4 , 4)
  284. ])
  285. r.drawLine(
  286. 4,
  287. TILE_PIXELS / 2,
  288. TILE_PIXELS - 4,
  289. TILE_PIXELS / 2
  290. )
  291. r.setLineWidth(1)
  292. def toggle(self, env, pos):
  293. # Replace the box by its contents
  294. env.grid.set(*pos, self.contains)
  295. return True
  296. class Grid:
  297. """
  298. Represent a grid and operations on it
  299. """
  300. # Static cache of pre-renderer tiles
  301. tile_cache = {}
  302. def __init__(self, width, height):
  303. assert width >= 3
  304. assert height >= 3
  305. self.width = width
  306. self.height = height
  307. self.grid = [None] * width * height
  308. def __contains__(self, key):
  309. if isinstance(key, WorldObj):
  310. for e in self.grid:
  311. if e is key:
  312. return True
  313. elif isinstance(key, tuple):
  314. for e in self.grid:
  315. if e is None:
  316. continue
  317. if (e.color, e.type) == key:
  318. return True
  319. if key[0] is None and key[1] == e.type:
  320. return True
  321. return False
  322. def __eq__(self, other):
  323. grid1 = self.encode()
  324. grid2 = other.encode()
  325. return np.array_equal(grid2, grid1)
  326. def __ne__(self, other):
  327. return not self == other
  328. def copy(self):
  329. from copy import deepcopy
  330. return deepcopy(self)
  331. def set(self, i, j, v):
  332. assert i >= 0 and i < self.width
  333. assert j >= 0 and j < self.height
  334. self.grid[j * self.width + i] = v
  335. def get(self, i, j):
  336. assert i >= 0 and i < self.width
  337. assert j >= 0 and j < self.height
  338. return self.grid[j * self.width + i]
  339. def horz_wall(self, x, y, length=None, obj_type=Wall):
  340. if length is None:
  341. length = self.width - x
  342. for i in range(0, length):
  343. self.set(x + i, y, obj_type())
  344. def vert_wall(self, x, y, length=None, obj_type=Wall):
  345. if length is None:
  346. length = self.height - y
  347. for j in range(0, length):
  348. self.set(x, y + j, obj_type())
  349. def wall_rect(self, x, y, w, h):
  350. self.horz_wall(x, y, w)
  351. self.horz_wall(x, y+h-1, w)
  352. self.vert_wall(x, y, h)
  353. self.vert_wall(x+w-1, y, h)
  354. def rotate_left(self):
  355. """
  356. Rotate the grid to the left (counter-clockwise)
  357. """
  358. grid = Grid(self.height, self.width)
  359. for i in range(self.width):
  360. for j in range(self.height):
  361. v = self.get(i, j)
  362. grid.set(j, grid.height - 1 - i, v)
  363. return grid
  364. def slice(self, topX, topY, width, height):
  365. """
  366. Get a subset of the grid
  367. """
  368. grid = Grid(width, height)
  369. for j in range(0, height):
  370. for i in range(0, width):
  371. x = topX + i
  372. y = topY + j
  373. if x >= 0 and x < self.width and \
  374. y >= 0 and y < self.height:
  375. v = self.get(x, y)
  376. else:
  377. v = Wall()
  378. grid.set(i, j, v)
  379. return grid
  380. @classmethod
  381. def render_tile(
  382. cls,
  383. obj,
  384. agent_dir=None,
  385. highlight=False,
  386. tile_size=TILE_PIXELS
  387. ):
  388. """
  389. Render a tile and cache the result
  390. """
  391. # Hash map lookup key for the cache
  392. key = (agent_dir, highlight, tile_size)
  393. key = obj.encode() + key if obj else key
  394. if key in cls.tile_cache:
  395. return cls.tile_cache[key]
  396. img = np.zeros(shape=(tile_size, tile_size, 3), dtype=np.uint8)
  397. # Draw the grid lines (top and left edges)
  398. fill_coords(img, point_in_rect(0, 0.031, 0, 1), (100, 100, 100))
  399. fill_coords(img, point_in_rect(0, 1, 0, 0.031), (100, 100, 100))
  400. if obj != None:
  401. obj.render(img)
  402. # Overlay the agent on top
  403. if agent_dir is not None:
  404. tri_fn = point_in_triangle(
  405. (0.12, 0.19),
  406. (0.87, 0.50),
  407. (0.12, 0.81),
  408. )
  409. # Rotate the agent based on its direction
  410. tri_fn = rotate_fn(tri_fn, cx=0.5, cy=0.5, theta=0.5*math.pi*agent_dir)
  411. fill_coords(img, tri_fn, (255, 0, 0))
  412. # Highlight the cell if needed
  413. if highlight:
  414. highlight_img(img)
  415. # Cache the rendered tile
  416. cls.tile_cache[key] = img
  417. return img
  418. def render(
  419. self,
  420. tile_size,
  421. agent_pos=None,
  422. agent_dir=None,
  423. highlight_mask=None
  424. ):
  425. """
  426. Render this grid at a given scale
  427. :param r: target renderer object
  428. :param tile_size: tile size in pixels
  429. """
  430. if highlight_mask is None:
  431. highlight_mask = np.zeros(shape=(self.width, self.height), dtype=np.bool)
  432. # Compute the total grid size
  433. width_px = self.width * TILE_PIXELS
  434. height_px = self.height * TILE_PIXELS
  435. img = np.zeros(shape=(height_px, width_px, 3), dtype=np.uint8)
  436. # Render the grid
  437. for j in range(0, self.height):
  438. for i in range(0, self.width):
  439. cell = self.get(i, j)
  440. agent_here = np.array_equal(agent_pos, (i, j))
  441. tile_img = Grid.render_tile(
  442. cell,
  443. agent_dir=agent_dir if agent_here else None,
  444. highlight=highlight_mask[i, j],
  445. tile_size=tile_size
  446. )
  447. ymin = j * tile_size
  448. ymax = (j+1) * tile_size
  449. xmin = i * tile_size
  450. xmax = (i+1) * tile_size
  451. img[ymin:ymax, xmin:xmax, :] = tile_img
  452. return img
  453. def encode(self, vis_mask=None):
  454. """
  455. Produce a compact numpy encoding of the grid
  456. """
  457. if vis_mask is None:
  458. vis_mask = np.ones((self.width, self.height), dtype=bool)
  459. array = np.zeros((self.width, self.height, 3), dtype='uint8')
  460. for i in range(self.width):
  461. for j in range(self.height):
  462. if vis_mask[i, j]:
  463. v = self.get(i, j)
  464. if v is None:
  465. array[i, j, 0] = OBJECT_TO_IDX['empty']
  466. array[i, j, 1] = 0
  467. array[i, j, 2] = 0
  468. else:
  469. array[i, j, :] = v.encode()
  470. return array
  471. @staticmethod
  472. def decode(array):
  473. """
  474. Decode an array grid encoding back into a grid
  475. """
  476. width, height, channels = array.shape
  477. assert channels == 3
  478. grid = Grid(width, height)
  479. for i in range(width):
  480. for j in range(height):
  481. type_idx, color_idx, state = array[i, j]
  482. v = WorldObj.decode(type_idx, color_idx, state)
  483. grid.set(i, j, v)
  484. return grid
  485. def process_vis(grid, agent_pos):
  486. mask = np.zeros(shape=(grid.width, grid.height), dtype=np.bool)
  487. mask[agent_pos[0], agent_pos[1]] = True
  488. for j in reversed(range(0, grid.height)):
  489. for i in range(0, grid.width-1):
  490. if not mask[i, j]:
  491. continue
  492. cell = grid.get(i, j)
  493. if cell and not cell.see_behind():
  494. continue
  495. mask[i+1, j] = True
  496. if j > 0:
  497. mask[i+1, j-1] = True
  498. mask[i, j-1] = True
  499. for i in reversed(range(1, grid.width)):
  500. if not mask[i, j]:
  501. continue
  502. cell = grid.get(i, j)
  503. if cell and not cell.see_behind():
  504. continue
  505. mask[i-1, j] = True
  506. if j > 0:
  507. mask[i-1, j-1] = True
  508. mask[i, j-1] = True
  509. for j in range(0, grid.height):
  510. for i in range(0, grid.width):
  511. if not mask[i, j]:
  512. grid.set(i, j, None)
  513. return mask
  514. class MiniGridEnv(gym.Env):
  515. """
  516. 2D grid world game environment
  517. """
  518. metadata = {
  519. 'render.modes': ['human', 'rgb_array', 'pixmap'],
  520. 'video.frames_per_second' : 10
  521. }
  522. # Enumeration of possible actions
  523. class Actions(IntEnum):
  524. # Turn left, turn right, move forward
  525. left = 0
  526. right = 1
  527. forward = 2
  528. # Pick up an object
  529. pickup = 3
  530. # Drop an object
  531. drop = 4
  532. # Toggle/activate an object
  533. toggle = 5
  534. # Done completing task
  535. done = 6
  536. def __init__(
  537. self,
  538. grid_size=None,
  539. width=None,
  540. height=None,
  541. max_steps=100,
  542. see_through_walls=False,
  543. seed=1337,
  544. agent_view_size=7
  545. ):
  546. # Can't set both grid_size and width/height
  547. if grid_size:
  548. assert width == None and height == None
  549. width = grid_size
  550. height = grid_size
  551. # Action enumeration for this environment
  552. self.actions = MiniGridEnv.Actions
  553. # Actions are discrete integer values
  554. self.action_space = spaces.Discrete(len(self.actions))
  555. # Number of cells (width and height) in the agent view
  556. self.agent_view_size = agent_view_size
  557. # Observations are dictionaries containing an
  558. # encoding of the grid and a textual 'mission' string
  559. self.observation_space = spaces.Box(
  560. low=0,
  561. high=255,
  562. shape=(self.agent_view_size, self.agent_view_size, 3),
  563. dtype='uint8'
  564. )
  565. self.observation_space = spaces.Dict({
  566. 'image': self.observation_space
  567. })
  568. # Range of possible rewards
  569. self.reward_range = (0, 1)
  570. # Renderer object used to render the whole grid (full-scale)
  571. self.grid_render = None
  572. # Renderer used to render observations (small-scale agent view)
  573. self.obs_render = None
  574. # Environment configuration
  575. self.width = width
  576. self.height = height
  577. self.max_steps = max_steps
  578. self.see_through_walls = see_through_walls
  579. # Current position and direction of the agent
  580. self.agent_pos = None
  581. self.agent_dir = None
  582. # Initialize the RNG
  583. self.seed(seed=seed)
  584. # Initialize the state
  585. self.reset()
  586. def reset(self):
  587. # Current position and direction of the agent
  588. self.agent_pos = None
  589. self.agent_dir = None
  590. # Generate a new random grid at the start of each episode
  591. # To keep the same grid for each episode, call env.seed() with
  592. # the same seed before calling env.reset()
  593. self._gen_grid(self.width, self.height)
  594. # These fields should be defined by _gen_grid
  595. assert self.agent_pos is not None
  596. assert self.agent_dir is not None
  597. # Check that the agent doesn't overlap with an object
  598. start_cell = self.grid.get(*self.agent_pos)
  599. assert start_cell is None or start_cell.can_overlap()
  600. # Item picked up, being carried, initially nothing
  601. self.carrying = None
  602. # Step count since episode start
  603. self.step_count = 0
  604. # Return first observation
  605. obs = self.gen_obs()
  606. return obs
  607. def seed(self, seed=1337):
  608. # Seed the random number generator
  609. self.np_random, _ = seeding.np_random(seed)
  610. return [seed]
  611. @property
  612. def steps_remaining(self):
  613. return self.max_steps - self.step_count
  614. def __str__(self):
  615. """
  616. Produce a pretty string of the environment's grid along with the agent.
  617. A grid cell is represented by 2-character string, the first one for
  618. the object and the second one for the color.
  619. """
  620. # Map of object types to short string
  621. OBJECT_TO_STR = {
  622. 'wall' : 'W',
  623. 'floor' : 'F',
  624. 'door' : 'D',
  625. 'key' : 'K',
  626. 'ball' : 'A',
  627. 'box' : 'B',
  628. 'goal' : 'G',
  629. 'lava' : 'V',
  630. }
  631. # Short string for opened door
  632. OPENDED_DOOR_IDS = '_'
  633. # Map agent's direction to short string
  634. AGENT_DIR_TO_STR = {
  635. 0: '>',
  636. 1: 'V',
  637. 2: '<',
  638. 3: '^'
  639. }
  640. str = ''
  641. for j in range(self.grid.height):
  642. for i in range(self.grid.width):
  643. if i == self.agent_pos[0] and j == self.agent_pos[1]:
  644. str += 2 * AGENT_DIR_TO_STR[self.agent_dir]
  645. continue
  646. c = self.grid.get(i, j)
  647. if c == None:
  648. str += ' '
  649. continue
  650. if c.type == 'door':
  651. if c.is_open:
  652. str += '__'
  653. elif c.is_locked:
  654. str += 'L' + c.color[0].upper()
  655. else:
  656. str += 'D' + c.color[0].upper()
  657. continue
  658. str += OBJECT_TO_STR[c.type] + c.color[0].upper()
  659. if j < self.grid.height - 1:
  660. str += '\n'
  661. return str
  662. def _gen_grid(self, width, height):
  663. assert False, "_gen_grid needs to be implemented by each environment"
  664. def _reward(self):
  665. """
  666. Compute the reward to be given upon success
  667. """
  668. return 1 - 0.9 * (self.step_count / self.max_steps)
  669. def _rand_int(self, low, high):
  670. """
  671. Generate random integer in [low,high[
  672. """
  673. return self.np_random.randint(low, high)
  674. def _rand_float(self, low, high):
  675. """
  676. Generate random float in [low,high[
  677. """
  678. return self.np_random.uniform(low, high)
  679. def _rand_bool(self):
  680. """
  681. Generate random boolean value
  682. """
  683. return (self.np_random.randint(0, 2) == 0)
  684. def _rand_elem(self, iterable):
  685. """
  686. Pick a random element in a list
  687. """
  688. lst = list(iterable)
  689. idx = self._rand_int(0, len(lst))
  690. return lst[idx]
  691. def _rand_subset(self, iterable, num_elems):
  692. """
  693. Sample a random subset of distinct elements of a list
  694. """
  695. lst = list(iterable)
  696. assert num_elems <= len(lst)
  697. out = []
  698. while len(out) < num_elems:
  699. elem = self._rand_elem(lst)
  700. lst.remove(elem)
  701. out.append(elem)
  702. return out
  703. def _rand_color(self):
  704. """
  705. Generate a random color name (string)
  706. """
  707. return self._rand_elem(COLOR_NAMES)
  708. def _rand_pos(self, xLow, xHigh, yLow, yHigh):
  709. """
  710. Generate a random (x,y) position tuple
  711. """
  712. return (
  713. self.np_random.randint(xLow, xHigh),
  714. self.np_random.randint(yLow, yHigh)
  715. )
  716. def place_obj(self,
  717. obj,
  718. top=None,
  719. size=None,
  720. reject_fn=None,
  721. max_tries=math.inf
  722. ):
  723. """
  724. Place an object at an empty position in the grid
  725. :param top: top-left position of the rectangle where to place
  726. :param size: size of the rectangle where to place
  727. :param reject_fn: function to filter out potential positions
  728. """
  729. if top is None:
  730. top = (0, 0)
  731. else:
  732. top = (max(top[0], 0), max(top[1], 0))
  733. if size is None:
  734. size = (self.grid.width, self.grid.height)
  735. num_tries = 0
  736. while True:
  737. # This is to handle with rare cases where rejection sampling
  738. # gets stuck in an infinite loop
  739. if num_tries > max_tries:
  740. raise RecursionError('rejection sampling failed in place_obj')
  741. num_tries += 1
  742. pos = np.array((
  743. self._rand_int(top[0], min(top[0] + size[0], self.grid.width)),
  744. self._rand_int(top[1], min(top[1] + size[1], self.grid.height))
  745. ))
  746. # Don't place the object on top of another object
  747. if self.grid.get(*pos) != None:
  748. continue
  749. # Don't place the object where the agent is
  750. if np.array_equal(pos, self.agent_pos):
  751. continue
  752. # Check if there is a filtering criterion
  753. if reject_fn and reject_fn(self, pos):
  754. continue
  755. break
  756. self.grid.set(*pos, obj)
  757. if obj is not None:
  758. obj.init_pos = pos
  759. obj.cur_pos = pos
  760. return pos
  761. def put_obj(self, obj, i, j):
  762. """
  763. Put an object at a specific position in the grid
  764. """
  765. self.grid.set(i, j, obj)
  766. obj.init_pos = (i, j)
  767. obj.cur_pos = (i, j)
  768. def place_agent(
  769. self,
  770. top=None,
  771. size=None,
  772. rand_dir=True,
  773. max_tries=math.inf
  774. ):
  775. """
  776. Set the agent's starting point at an empty position in the grid
  777. """
  778. self.agent_pos = None
  779. pos = self.place_obj(None, top, size, max_tries=max_tries)
  780. self.agent_pos = pos
  781. if rand_dir:
  782. self.agent_dir = self._rand_int(0, 4)
  783. return pos
  784. @property
  785. def dir_vec(self):
  786. """
  787. Get the direction vector for the agent, pointing in the direction
  788. of forward movement.
  789. """
  790. assert self.agent_dir >= 0 and self.agent_dir < 4
  791. return DIR_TO_VEC[self.agent_dir]
  792. @property
  793. def right_vec(self):
  794. """
  795. Get the vector pointing to the right of the agent.
  796. """
  797. dx, dy = self.dir_vec
  798. return np.array((-dy, dx))
  799. @property
  800. def front_pos(self):
  801. """
  802. Get the position of the cell that is right in front of the agent
  803. """
  804. return self.agent_pos + self.dir_vec
  805. def get_view_coords(self, i, j):
  806. """
  807. Translate and rotate absolute grid coordinates (i, j) into the
  808. agent's partially observable view (sub-grid). Note that the resulting
  809. coordinates may be negative or outside of the agent's view size.
  810. """
  811. ax, ay = self.agent_pos
  812. dx, dy = self.dir_vec
  813. rx, ry = self.right_vec
  814. # Compute the absolute coordinates of the top-left view corner
  815. sz = self.agent_view_size
  816. hs = self.agent_view_size // 2
  817. tx = ax + (dx * (sz-1)) - (rx * hs)
  818. ty = ay + (dy * (sz-1)) - (ry * hs)
  819. lx = i - tx
  820. ly = j - ty
  821. # Project the coordinates of the object relative to the top-left
  822. # corner onto the agent's own coordinate system
  823. vx = (rx*lx + ry*ly)
  824. vy = -(dx*lx + dy*ly)
  825. return vx, vy
  826. def get_view_exts(self):
  827. """
  828. Get the extents of the square set of tiles visible to the agent
  829. Note: the bottom extent indices are not included in the set
  830. """
  831. # Facing right
  832. if self.agent_dir == 0:
  833. topX = self.agent_pos[0]
  834. topY = self.agent_pos[1] - self.agent_view_size // 2
  835. # Facing down
  836. elif self.agent_dir == 1:
  837. topX = self.agent_pos[0] - self.agent_view_size // 2
  838. topY = self.agent_pos[1]
  839. # Facing left
  840. elif self.agent_dir == 2:
  841. topX = self.agent_pos[0] - self.agent_view_size + 1
  842. topY = self.agent_pos[1] - self.agent_view_size // 2
  843. # Facing up
  844. elif self.agent_dir == 3:
  845. topX = self.agent_pos[0] - self.agent_view_size // 2
  846. topY = self.agent_pos[1] - self.agent_view_size + 1
  847. else:
  848. assert False, "invalid agent direction"
  849. botX = topX + self.agent_view_size
  850. botY = topY + self.agent_view_size
  851. return (topX, topY, botX, botY)
  852. def relative_coords(self, x, y):
  853. """
  854. Check if a grid position belongs to the agent's field of view, and returns the corresponding coordinates
  855. """
  856. vx, vy = self.get_view_coords(x, y)
  857. if vx < 0 or vy < 0 or vx >= self.agent_view_size or vy >= self.agent_view_size:
  858. return None
  859. return vx, vy
  860. def in_view(self, x, y):
  861. """
  862. check if a grid position is visible to the agent
  863. """
  864. return self.relative_coords(x, y) is not None
  865. def agent_sees(self, x, y):
  866. """
  867. Check if a non-empty grid position is visible to the agent
  868. """
  869. coordinates = self.relative_coords(x, y)
  870. if coordinates is None:
  871. return False
  872. vx, vy = coordinates
  873. obs = self.gen_obs()
  874. obs_grid = Grid.decode(obs['image'])
  875. obs_cell = obs_grid.get(vx, vy)
  876. world_cell = self.grid.get(x, y)
  877. return obs_cell is not None and obs_cell.type == world_cell.type
  878. def step(self, action):
  879. self.step_count += 1
  880. reward = 0
  881. done = False
  882. # Get the position in front of the agent
  883. fwd_pos = self.front_pos
  884. # Get the contents of the cell in front of the agent
  885. fwd_cell = self.grid.get(*fwd_pos)
  886. # Rotate left
  887. if action == self.actions.left:
  888. self.agent_dir -= 1
  889. if self.agent_dir < 0:
  890. self.agent_dir += 4
  891. # Rotate right
  892. elif action == self.actions.right:
  893. self.agent_dir = (self.agent_dir + 1) % 4
  894. # Move forward
  895. elif action == self.actions.forward:
  896. if fwd_cell == None or fwd_cell.can_overlap():
  897. self.agent_pos = fwd_pos
  898. if fwd_cell != None and fwd_cell.type == 'goal':
  899. done = True
  900. reward = self._reward()
  901. if fwd_cell != None and fwd_cell.type == 'lava':
  902. done = True
  903. # Pick up an object
  904. elif action == self.actions.pickup:
  905. if fwd_cell and fwd_cell.can_pickup():
  906. if self.carrying is None:
  907. self.carrying = fwd_cell
  908. self.carrying.cur_pos = np.array([-1, -1])
  909. self.grid.set(*fwd_pos, None)
  910. # Drop an object
  911. elif action == self.actions.drop:
  912. if not fwd_cell and self.carrying:
  913. self.grid.set(*fwd_pos, self.carrying)
  914. self.carrying.cur_pos = fwd_pos
  915. self.carrying = None
  916. # Toggle/activate an object
  917. elif action == self.actions.toggle:
  918. if fwd_cell:
  919. fwd_cell.toggle(self, fwd_pos)
  920. # Done action (not used by default)
  921. elif action == self.actions.done:
  922. pass
  923. else:
  924. assert False, "unknown action"
  925. if self.step_count >= self.max_steps:
  926. done = True
  927. obs = self.gen_obs()
  928. return obs, reward, done, {}
  929. def gen_obs_grid(self):
  930. """
  931. Generate the sub-grid observed by the agent.
  932. This method also outputs a visibility mask telling us which grid
  933. cells the agent can actually see.
  934. """
  935. topX, topY, botX, botY = self.get_view_exts()
  936. grid = self.grid.slice(topX, topY, self.agent_view_size, self.agent_view_size)
  937. for i in range(self.agent_dir + 1):
  938. grid = grid.rotate_left()
  939. # Process occluders and visibility
  940. # Note that this incurs some performance cost
  941. if not self.see_through_walls:
  942. vis_mask = grid.process_vis(agent_pos=(self.agent_view_size // 2 , self.agent_view_size - 1))
  943. else:
  944. vis_mask = np.ones(shape=(grid.width, grid.height), dtype=np.bool)
  945. # Make it so the agent sees what it's carrying
  946. # We do this by placing the carried object at the agent's position
  947. # in the agent's partially observable view
  948. agent_pos = grid.width // 2, grid.height - 1
  949. if self.carrying:
  950. grid.set(*agent_pos, self.carrying)
  951. else:
  952. grid.set(*agent_pos, None)
  953. return grid, vis_mask
  954. def gen_obs(self):
  955. """
  956. Generate the agent's view (partially observable, low-resolution encoding)
  957. """
  958. grid, vis_mask = self.gen_obs_grid()
  959. # Encode the partially observable view into a numpy array
  960. image = grid.encode(vis_mask)
  961. assert hasattr(self, 'mission'), "environments must define a textual mission string"
  962. # Observations are dictionaries containing:
  963. # - an image (partially observable view of the environment)
  964. # - the agent's direction/orientation (acting as a compass)
  965. # - a textual mission string (instructions for the agent)
  966. obs = {
  967. 'image': image,
  968. 'direction': self.agent_dir,
  969. 'mission': self.mission
  970. }
  971. return obs
  972. def get_obs_render(self, obs, tile_size=TILE_PIXELS//2):
  973. """
  974. Render an agent observation for visualization
  975. """
  976. grid = Grid.decode(obs)
  977. # Render the whole grid
  978. img = grid.render(r, tile_size)
  979. assert False
  980. """
  981. # Draw the agent
  982. ratio = tile_size / TILE_PIXELS
  983. r.push()
  984. r.scale(ratio, ratio)
  985. r.translate(
  986. TILE_PIXELS * (0.5 + self.agent_view_size // 2),
  987. TILE_PIXELS * (self.agent_view_size - 0.5)
  988. )
  989. r.rotate(3 * 90)
  990. r.setLineColor(255, 0, 0)
  991. r.setColor(255, 0, 0)
  992. r.drawPolygon([
  993. (-12, 10),
  994. ( 12, 0),
  995. (-12, -10)
  996. ])
  997. r.pop()
  998. """
  999. return img
  1000. def render(self, mode='human', close=False, highlight=True, tile_size=TILE_PIXELS):
  1001. """
  1002. Render the whole-grid human view
  1003. """
  1004. """
  1005. if close:
  1006. if self.grid_render:
  1007. self.grid_render.close()
  1008. return
  1009. """
  1010. # Compute which cells are visible to the agent
  1011. _, vis_mask = self.gen_obs_grid()
  1012. # Compute the world coordinates of the bottom-left corner
  1013. # of the agent's view area
  1014. f_vec = self.dir_vec
  1015. r_vec = self.right_vec
  1016. top_left = self.agent_pos + f_vec * (self.agent_view_size-1) - r_vec * (self.agent_view_size // 2)
  1017. # Mask of which cells to highlight
  1018. highlight_mask = np.zeros(shape=(self.width, self.height), dtype=np.bool)
  1019. # For each cell in the visibility mask
  1020. for vis_j in range(0, self.agent_view_size):
  1021. for vis_i in range(0, self.agent_view_size):
  1022. # If this cell is not visible, don't highlight it
  1023. if not vis_mask[vis_i, vis_j]:
  1024. continue
  1025. # Compute the world coordinates of this cell
  1026. abs_i, abs_j = top_left - (f_vec * vis_j) + (r_vec * vis_i)
  1027. if abs_i < 0 or abs_i >= self.width:
  1028. continue
  1029. if abs_j < 0 or abs_j >= self.height:
  1030. continue
  1031. # Mark this cell to be highlighted
  1032. highlight_mask[abs_i, abs_j] = True
  1033. # Render the whole grid
  1034. img = self.grid.render(
  1035. tile_size,
  1036. self.agent_pos,
  1037. self.agent_dir,
  1038. highlight_mask=highlight_mask if highlight else None
  1039. )
  1040. return img