The source code and dockerfile for the GSW2024 AI Lab.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
This repo is archived. You can view files and clone it, but cannot push or open issues/pull-requests.

201 lines
5.8 KiB

2 months ago
  1. """
  2. Copied and adapted from https://github.com/mila-iqia/babyai.
  3. Levels described in the Baby AI ICLR 2019 submission, with the `Put Next` instruction.
  4. """
  5. from __future__ import annotations
  6. from minigrid.envs.babyai.core.roomgrid_level import RoomGridLevel
  7. from minigrid.envs.babyai.core.verifier import ObjDesc, PutNextInstr
  8. class PutNextLocal(RoomGridLevel):
  9. """
  10. ## Description
  11. Put an object next to another object, inside a single room
  12. with no doors, no distractors
  13. ## Mission Space
  14. "put the {color} {type} next to the {color} {type}"
  15. {color} is the color of the box. Can be "red", "green", "blue", "purple",
  16. "yellow" or "grey".
  17. {type} is the type of the object. Can be "ball", "box" or "key".
  18. ## Action Space
  19. | Num | Name | Action |
  20. |-----|--------------|-------------------|
  21. | 0 | left | Turn left |
  22. | 1 | right | Turn right |
  23. | 2 | forward | Move forward |
  24. | 3 | pickup | Pick up an object |
  25. | 4 | drop | Unused |
  26. | 5 | toggle | Unused |
  27. | 6 | done | Unused |
  28. ## Observation Encoding
  29. - Each tile is encoded as a 3 dimensional tuple:
  30. `(OBJECT_IDX, COLOR_IDX, STATE)`
  31. - `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
  32. [minigrid/minigrid.py](minigrid/minigrid.py)
  33. - `STATE` refers to the door state with 0=open, 1=closed and 2=locked
  34. ## Rewards
  35. A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
  36. ## Termination
  37. The episode ends if any one of the following conditions is met:
  38. 1. The agent finishes the instructed task.
  39. 2. Timeout (see `max_steps`).
  40. ## Registered Configurations
  41. - `BabyAI-PutNextLocal-v0`
  42. - `BabyAI-PutNextLocalS5N3-v0`
  43. - `BabyAI-PutNextLocalS6N4-v0``
  44. """
  45. def __init__(self, room_size=8, num_objs=8, **kwargs):
  46. self.num_objs = num_objs
  47. super().__init__(num_rows=1, num_cols=1, room_size=room_size, **kwargs)
  48. def gen_mission(self):
  49. self.place_agent()
  50. objs = self.add_distractors(num_distractors=self.num_objs, all_unique=True)
  51. self.check_objs_reachable()
  52. o1, o2 = self._rand_subset(objs, 2)
  53. self.instrs = PutNextInstr(
  54. ObjDesc(o1.type, o1.color), ObjDesc(o2.type, o2.color)
  55. )
  56. class PutNext(RoomGridLevel):
  57. """
  58. ## Description
  59. Task of the form: move the A next to the B and the C next to the D.
  60. This task is structured to have a very large number of possible
  61. instructions.
  62. ## Mission Space
  63. "put the {color} {type} next to the {color} {type}"
  64. {color} is the color of the box. Can be "red", "green", "blue", "purple",
  65. "yellow" or "grey".
  66. {type} is the type of the object. Can be "ball", "box" or "key".
  67. ## Action Space
  68. | Num | Name | Action |
  69. |-----|--------------|-------------------|
  70. | 0 | left | Turn left |
  71. | 1 | right | Turn right |
  72. | 2 | forward | Move forward |
  73. | 3 | pickup | Pick up an object |
  74. | 4 | drop | Unused |
  75. | 5 | toggle | Unused |
  76. | 6 | done | Unused |
  77. ## Observation Encoding
  78. - Each tile is encoded as a 3 dimensional tuple:
  79. `(OBJECT_IDX, COLOR_IDX, STATE)`
  80. - `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
  81. [minigrid/minigrid.py](minigrid/minigrid.py)
  82. - `STATE` refers to the door state with 0=open, 1=closed and 2=locked
  83. ## Rewards
  84. A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
  85. ## Termination
  86. The episode ends if any one of the following conditions is met:
  87. 1. The agent finishes the instructed task.
  88. 2. Timeout (see `max_steps`).
  89. ## Registered Configurations
  90. - `BabyAI-PutNextS4N1-v0`
  91. - `BabyAI-PutNextS5N2-v0`
  92. - `BabyAI-PutNextS5N1-v0`
  93. - `BabyAI-PutNextS6N3-v0`
  94. - `BabyAI-PutNextS7N4-v0`
  95. - `BabyAI-PutNextS5N2Carrying-v0`
  96. - `BabyAI-PutNextS6N3Carrying-v0`
  97. - `BabyAI-PutNextS7N4Carrying-v0`
  98. ## Additional Notes
  99. The BabyAI bot is unable to solve the bonus PutNextCarrying configurations.
  100. """
  101. def __init__(
  102. self,
  103. room_size,
  104. objs_per_room,
  105. start_carrying=False,
  106. max_steps: int | None = None,
  107. **kwargs,
  108. ):
  109. assert room_size >= 4
  110. assert objs_per_room <= 9
  111. self.objs_per_room = objs_per_room
  112. self.start_carrying = start_carrying
  113. if max_steps is None:
  114. max_steps = 8 * room_size**2
  115. super().__init__(
  116. num_rows=1, num_cols=2, room_size=room_size, max_steps=max_steps, **kwargs
  117. )
  118. def gen_mission(self):
  119. self.place_agent(0, 0)
  120. # Add objects to both the left and right rooms
  121. # so that we know that we have two non-adjacent set of objects
  122. objs_l = self.add_distractors(0, 0, self.objs_per_room)
  123. objs_r = self.add_distractors(1, 0, self.objs_per_room)
  124. # Remove the wall between the two rooms
  125. self.remove_wall(0, 0, 0)
  126. # Select objects from both subsets
  127. a = self._rand_elem(objs_l)
  128. b = self._rand_elem(objs_r)
  129. # Randomly flip the object to be moved
  130. if self._rand_bool():
  131. t = a
  132. a = b
  133. b = t
  134. self.obj_a = a
  135. self.instrs = PutNextInstr(ObjDesc(a.type, a.color), ObjDesc(b.type, b.color))
  136. def reset(self, **kwargs):
  137. obs = super().reset(**kwargs)
  138. # If the agent starts off carrying the object
  139. if self.start_carrying:
  140. assert self.obj_a.init_pos is not None
  141. self.grid.set(*self.obj_a.init_pos, None)
  142. self.carrying = self.obj_a
  143. return obs