fourroomqa.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. from gym_minigrid.minigrid import *
  2. from gym_minigrid.register import register
  3. """
  4. class Room:
  5. def __init__(self,
  6. top,
  7. size,
  8. entryDoorPos,
  9. exitDoorPos
  10. ):
  11. self.top = top
  12. self.size = size
  13. self.entryDoorPos = entryDoorPos
  14. self.exitDoorPos = exitDoorPos
  15. """
  16. class FourRoomQAEnv(MiniGridEnv):
  17. """
  18. Environment to experiment with embodied question answering
  19. https://arxiv.org/abs/1711.11543
  20. """
  21. # Enumeration of possible actions
  22. class Actions(IntEnum):
  23. left = 0
  24. right = 1
  25. forward = 2
  26. toggle = 3
  27. say = 4
  28. def __init__(self, size=16):
  29. assert size >= 8
  30. super(FourRoomQAEnv, self).__init__(gridSize=size, maxSteps=8*size)
  31. # Action enumeration for this environment
  32. self.actions = MiniGridEnv.Actions
  33. # TODO: dictionary action space, to include answer sentence?
  34. # Actions are discrete integer values
  35. self.action_space = spaces.Discrete(len(self.actions))
  36. def _genGrid(self, width, height):
  37. grid = super(FourRoomQAEnv, self)._genGrid(width, height)
  38. # TODO: work in progress
  39. """
  40. # Create a vertical splitting wall
  41. splitIdx = self._randInt(2, gridSz-3)
  42. for i in range(0, gridSz):
  43. grid.set(splitIdx, i, Wall())
  44. # Place a door in the wall
  45. doorIdx = self._randInt(1, gridSz-2)
  46. grid.set(splitIdx, doorIdx, Door('yellow'))
  47. # Place a key on the left side
  48. #keyIdx = self._randInt(1 + gridSz // 2, gridSz-2)
  49. keyIdx = gridSz-2
  50. grid.set(1, keyIdx, Key('yellow'))
  51. """
  52. return grid
  53. register(
  54. id='MiniGrid-FourRoomQA-v0',
  55. entry_point='gym_minigrid.envs:FourRoomQAEnv'
  56. )