fetch.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127
  1. from gym_minigrid.minigrid import *
  2. from gym_minigrid.register import register
  3. class FetchEnv(MiniGridEnv):
  4. """
  5. Environment in which the agent has to fetch a random object
  6. named using English text strings
  7. """
  8. def __init__(
  9. self,
  10. size=8,
  11. numObjs=3
  12. ):
  13. self.numObjs = numObjs
  14. super().__init__(gridSize=size, maxSteps=5*size)
  15. self.reward_range = (-1000, 1000)
  16. def _genGrid(self, width, height):
  17. assert width == height
  18. gridSz = width
  19. # Create a grid surrounded by walls
  20. grid = Grid(width, height)
  21. for i in range(0, width):
  22. grid.set(i, 0, Wall())
  23. grid.set(i, height-1, Wall())
  24. for j in range(0, height):
  25. grid.set(0, j, Wall())
  26. grid.set(width-1, j, Wall())
  27. types = ['key', 'ball']
  28. colors = list(COLORS.keys())
  29. objs = []
  30. # For each object to be generated
  31. for i in range(0, self.numObjs):
  32. objType = self._randElem(types)
  33. objColor = self._randElem(colors)
  34. if objType == 'key':
  35. obj = Key(objColor)
  36. elif objType == 'ball':
  37. obj = Ball(objColor)
  38. while True:
  39. pos = (
  40. self._randInt(1, gridSz - 1),
  41. self._randInt(1, gridSz - 1)
  42. )
  43. if pos != self.startPos:
  44. grid.set(*pos, obj)
  45. break
  46. objs.append(obj)
  47. # Choose a random object to be picked up
  48. target = objs[self._randInt(0, len(objs))]
  49. self.targetType = target.type
  50. self.targetColor = target.color
  51. descStr = '%s %s' % (self.targetColor, self.targetType)
  52. # Generate the mission string
  53. idx = self._randInt(0, 5)
  54. if idx == 0:
  55. self.mission = 'get a %s' % descStr
  56. elif idx == 1:
  57. self.mission = 'go get a %s' % descStr
  58. elif idx == 2:
  59. self.mission = 'fetch a %s' % descStr
  60. elif idx == 3:
  61. self.mission = 'go fetch a %s' % descStr
  62. elif idx == 4:
  63. self.mission = 'you must fetch a %s' % descStr
  64. assert hasattr(self, 'mission')
  65. return grid
  66. def _observation(self, obs):
  67. """
  68. Encode observations
  69. """
  70. obs = {
  71. 'image': obs,
  72. 'mission': self.mission
  73. }
  74. return obs
  75. def _reset(self):
  76. obs = MiniGridEnv._reset(self)
  77. return self._observation(obs)
  78. def _step(self, action):
  79. obs, reward, done, info = MiniGridEnv._step(self, action)
  80. if self.carrying:
  81. if self.carrying.color == self.targetColor and \
  82. self.carrying.type == self.targetType:
  83. reward = 1000 - self.stepCount
  84. done = True
  85. else:
  86. reward = -1000
  87. done = True
  88. obs = self._observation(obs)
  89. return obs, reward, done, info
  90. class FetchEnv5x5N2(FetchEnv):
  91. def __init__(self):
  92. super().__init__(size=5, numObjs=2)
  93. register(
  94. id='MiniGrid-Fetch-5x5-N2-v0',
  95. entry_point='gym_minigrid.envs:FetchEnv5x5N2'
  96. )
  97. register(
  98. id='MiniGrid-Fetch-8x8-N3-v0',
  99. entry_point='gym_minigrid.envs:FetchEnv'
  100. )