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.

184 lines
6.1 KiB

2 months ago
  1. from __future__ import annotations
  2. import itertools as itt
  3. import numpy as np
  4. from minigrid.core.grid import Grid
  5. from minigrid.core.mission import MissionSpace
  6. from minigrid.core.world_object import Goal, Lava
  7. from minigrid.minigrid_env import MiniGridEnv
  8. class CrossingEnv(MiniGridEnv):
  9. """
  10. ## Description
  11. Depending on the `obstacle_type` parameter:
  12. - `Lava` - The agent has to reach the green goal square on the other corner
  13. of the room while avoiding rivers of deadly lava which terminate the
  14. episode in failure. Each lava stream runs across the room either
  15. horizontally or vertically, and has a single crossing point which can be
  16. safely used; Luckily, a path to the goal is guaranteed to exist. This
  17. environment is useful for studying safety and safe exploration.
  18. - otherwise - Similar to the `LavaCrossing` environment, the agent has to
  19. reach the green goal square on the other corner of the room, however
  20. lava is replaced by walls. This MDP is therefore much easier and maybe
  21. useful for quickly testing your algorithms.
  22. ## Mission Space
  23. Depending on the `obstacle_type` parameter:
  24. - `Lava` - "avoid the lava and get to the green goal square"
  25. - otherwise - "find the opening and get to the green goal square"
  26. ## Action Space
  27. | Num | Name | Action |
  28. |-----|--------------|--------------|
  29. | 0 | left | Turn left |
  30. | 1 | right | Turn right |
  31. | 2 | forward | Move forward |
  32. | 3 | pickup | Unused |
  33. | 4 | drop | Unused |
  34. | 5 | toggle | Unused |
  35. | 6 | done | Unused |
  36. ## Observation Encoding
  37. - Each tile is encoded as a 3 dimensional tuple:
  38. `(OBJECT_IDX, COLOR_IDX, STATE)`
  39. - `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
  40. [minigrid/minigrid.py](minigrid/minigrid.py)
  41. - `STATE` refers to the door state with 0=open, 1=closed and 2=locked
  42. ## Rewards
  43. A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
  44. ## Termination
  45. The episode ends if any one of the following conditions is met:
  46. 1. The agent reaches the goal.
  47. 2. The agent falls into lava.
  48. 3. Timeout (see `max_steps`).
  49. ## Registered Configurations
  50. S: size of the map SxS.
  51. N: number of valid crossings across lava or walls from the starting position
  52. to the goal
  53. - `Lava` :
  54. - `MiniGrid-LavaCrossingS9N1-v0`
  55. - `MiniGrid-LavaCrossingS9N2-v0`
  56. - `MiniGrid-LavaCrossingS9N3-v0`
  57. - `MiniGrid-LavaCrossingS11N5-v0`
  58. - otherwise :
  59. - `MiniGrid-SimpleCrossingS9N1-v0`
  60. - `MiniGrid-SimpleCrossingS9N2-v0`
  61. - `MiniGrid-SimpleCrossingS9N3-v0`
  62. - `MiniGrid-SimpleCrossingS11N5-v0`
  63. """
  64. def __init__(
  65. self,
  66. size=9,
  67. num_crossings=1,
  68. obstacle_type=Lava,
  69. max_steps: int | None = None,
  70. **kwargs,
  71. ):
  72. self.num_crossings = num_crossings
  73. self.obstacle_type = obstacle_type
  74. if obstacle_type == Lava:
  75. mission_space = MissionSpace(mission_func=self._gen_mission_lava)
  76. else:
  77. mission_space = MissionSpace(mission_func=self._gen_mission)
  78. if max_steps is None:
  79. max_steps = 4 * size**2
  80. super().__init__(
  81. mission_space=mission_space,
  82. grid_size=size,
  83. see_through_walls=False, # Set this to True for maximum speed
  84. max_steps=max_steps,
  85. **kwargs,
  86. )
  87. @staticmethod
  88. def _gen_mission_lava():
  89. return "avoid the lava and get to the green goal square"
  90. @staticmethod
  91. def _gen_mission():
  92. return "find the opening and get to the green goal square"
  93. def _gen_grid(self, width, height):
  94. assert width % 2 == 1 and height % 2 == 1 # odd size
  95. # Create an empty grid
  96. self.grid = Grid(width, height)
  97. # Generate the surrounding walls
  98. self.grid.wall_rect(0, 0, width, height)
  99. # Place the agent in the top-left corner
  100. self.agent_pos = np.array((1, 1))
  101. self.agent_dir = 0
  102. # Place a goal square in the bottom-right corner
  103. self.put_obj(Goal(), width - 2, height - 2)
  104. # Place obstacles (lava or walls)
  105. v, h = object(), object() # singleton `vertical` and `horizontal` objects
  106. # Lava rivers or walls specified by direction and position in grid
  107. rivers = [(v, i) for i in range(2, height - 2, 2)]
  108. rivers += [(h, j) for j in range(2, width - 2, 2)]
  109. self.np_random.shuffle(rivers)
  110. rivers = rivers[: self.num_crossings] # sample random rivers
  111. rivers_v = sorted(pos for direction, pos in rivers if direction is v)
  112. rivers_h = sorted(pos for direction, pos in rivers if direction is h)
  113. obstacle_pos = itt.chain(
  114. itt.product(range(1, width - 1), rivers_h),
  115. itt.product(rivers_v, range(1, height - 1)),
  116. )
  117. for i, j in obstacle_pos:
  118. self.put_obj(self.obstacle_type(), i, j)
  119. # Sample path to goal
  120. path = [h] * len(rivers_v) + [v] * len(rivers_h)
  121. self.np_random.shuffle(path)
  122. # Create openings
  123. limits_v = [0] + rivers_v + [height - 1]
  124. limits_h = [0] + rivers_h + [width - 1]
  125. room_i, room_j = 0, 0
  126. for direction in path:
  127. if direction is h:
  128. i = limits_v[room_i + 1]
  129. j = self.np_random.choice(
  130. range(limits_h[room_j] + 1, limits_h[room_j + 1])
  131. )
  132. room_i += 1
  133. elif direction is v:
  134. i = self.np_random.choice(
  135. range(limits_v[room_i] + 1, limits_v[room_i + 1])
  136. )
  137. j = limits_h[room_j + 1]
  138. room_j += 1
  139. else:
  140. assert False
  141. self.grid.set(i, j, None)
  142. self.mission = (
  143. "avoid the lava and get to the green goal square"
  144. if self.obstacle_type == Lava
  145. else "find the opening and get to the green goal square"
  146. )