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.
361 lines
10 KiB
361 lines
10 KiB
"""
|
|
Copied and adapted from https://github.com/mila-iqia/babyai.
|
|
Levels described in the Baby AI ICLR 2019 submission, with the `Pick up` instruction.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from minigrid.envs.babyai.core.levelgen import LevelGen
|
|
from minigrid.envs.babyai.core.roomgrid_level import RejectSampling, RoomGridLevel
|
|
from minigrid.envs.babyai.core.verifier import ObjDesc, PickupInstr
|
|
|
|
|
|
class Pickup(RoomGridLevel):
|
|
"""
|
|
|
|
## Description
|
|
|
|
Pick up an object, the object may be in another room.
|
|
|
|
## Mission Space
|
|
|
|
"pick up a {color} {type}"
|
|
|
|
{color} is the color of the box. Can be "red", "green", "blue", "purple",
|
|
"yellow" or "grey".
|
|
|
|
{type} is the type of the object. Can be "ball", "box" or "key".
|
|
|
|
## Action Space
|
|
|
|
| Num | Name | Action |
|
|
|-----|--------------|-------------------|
|
|
| 0 | left | Turn left |
|
|
| 1 | right | Turn right |
|
|
| 2 | forward | Move forward |
|
|
| 3 | pickup | Pick up an object |
|
|
| 4 | drop | Unused |
|
|
| 5 | toggle | Unused |
|
|
| 6 | done | Unused |
|
|
|
|
## Observation Encoding
|
|
|
|
- Each tile is encoded as a 3 dimensional tuple:
|
|
`(OBJECT_IDX, COLOR_IDX, STATE)`
|
|
- `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
|
|
[minigrid/minigrid.py](minigrid/minigrid.py)
|
|
- `STATE` refers to the door state with 0=open, 1=closed and 2=locked
|
|
|
|
## Rewards
|
|
|
|
A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
|
|
|
|
## Termination
|
|
|
|
The episode ends if any one of the following conditions is met:
|
|
|
|
1. The agent picks up the object.
|
|
2. Timeout (see `max_steps`).
|
|
|
|
## Registered Configurations
|
|
|
|
- `BabyAI-Pickup-v0`
|
|
|
|
"""
|
|
|
|
def gen_mission(self):
|
|
self.place_agent()
|
|
self.connect_all()
|
|
objs = self.add_distractors(num_distractors=18, all_unique=False)
|
|
self.check_objs_reachable()
|
|
obj = self._rand_elem(objs)
|
|
self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
|
|
|
|
|
|
class UnblockPickup(RoomGridLevel):
|
|
"""
|
|
|
|
## Description
|
|
|
|
Pick up an object, the object may be in another room. The path may
|
|
be blocked by one or more obstructors.
|
|
|
|
## Mission Space
|
|
|
|
"pick up a/the {color} {type}"
|
|
|
|
{color} is the color of the box. Can be "red", "green", "blue", "purple",
|
|
"yellow" or "grey".
|
|
|
|
{type} is the type of the object. Can be "ball", "box" or "key".
|
|
|
|
## Action Space
|
|
|
|
| Num | Name | Action |
|
|
|-----|--------------|-------------------|
|
|
| 0 | left | Turn left |
|
|
| 1 | right | Turn right |
|
|
| 2 | forward | Move forward |
|
|
| 3 | pickup | Pick up an object |
|
|
| 4 | drop | Unused |
|
|
| 5 | toggle | Unused |
|
|
| 6 | done | Unused |
|
|
|
|
## Observation Encoding
|
|
|
|
- Each tile is encoded as a 3 dimensional tuple:
|
|
`(OBJECT_IDX, COLOR_IDX, STATE)`
|
|
- `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
|
|
[minigrid/minigrid.py](minigrid/minigrid.py)
|
|
- `STATE` refers to the door state with 0=open, 1=closed and 2=locked
|
|
|
|
## Rewards
|
|
|
|
A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
|
|
|
|
## Termination
|
|
|
|
The episode ends if any one of the following conditions is met:
|
|
|
|
1. The agent picks up the object.
|
|
2. Timeout (see `max_steps`).
|
|
|
|
## Registered Configurations
|
|
|
|
- `BabyAI-UnblockPickup-v0`
|
|
|
|
"""
|
|
|
|
def gen_mission(self):
|
|
self.place_agent()
|
|
self.connect_all()
|
|
objs = self.add_distractors(num_distractors=20, all_unique=False)
|
|
|
|
# Ensure that at least one object is not reachable without unblocking
|
|
# Note: the selected object will still be reachable most of the time
|
|
if self.check_objs_reachable(raise_exc=False):
|
|
raise RejectSampling("all objects reachable")
|
|
|
|
obj = self._rand_elem(objs)
|
|
self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
|
|
|
|
|
|
class PickupLoc(LevelGen):
|
|
"""
|
|
|
|
## Description
|
|
|
|
Pick up an object which may be described using its location. This is a
|
|
single room environment.
|
|
|
|
Competencies: PickUp, Loc. No unblocking.
|
|
|
|
## Mission Space
|
|
|
|
"pick up the {color} {type}"
|
|
|
|
{color} is the color of the box. Can be "red", "green", "blue", "purple",
|
|
"yellow" or "grey".
|
|
|
|
{type} is the type of the object. Can be "ball", "box" or "key".
|
|
|
|
## Action Space
|
|
|
|
| Num | Name | Action |
|
|
|-----|--------------|-------------------|
|
|
| 0 | left | Turn left |
|
|
| 1 | right | Turn right |
|
|
| 2 | forward | Move forward |
|
|
| 3 | pickup | Pick up an object |
|
|
| 4 | drop | Unused |
|
|
| 5 | toggle | Unused |
|
|
| 6 | done | Unused |
|
|
|
|
## Observation Encoding
|
|
|
|
- Each tile is encoded as a 3 dimensional tuple:
|
|
`(OBJECT_IDX, COLOR_IDX, STATE)`
|
|
- `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
|
|
[minigrid/minigrid.py](minigrid/minigrid.py)
|
|
- `STATE` refers to the door state with 0=open, 1=closed and 2=locked
|
|
|
|
## Rewards
|
|
|
|
A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
|
|
|
|
## Termination
|
|
|
|
The episode ends if any one of the following conditions is met:
|
|
|
|
1. The agent picks up the object.
|
|
2. Timeout (see `max_steps`).
|
|
|
|
## Registered Configurations
|
|
|
|
- `BabyAI-PickupLoc-v0`
|
|
|
|
"""
|
|
|
|
def __init__(self, **kwargs):
|
|
# We add many distractors to increase the probability
|
|
# of ambiguous locations within the same room
|
|
super().__init__(
|
|
action_kinds=["pickup"],
|
|
instr_kinds=["action"],
|
|
num_rows=1,
|
|
num_cols=1,
|
|
num_dists=8,
|
|
locked_room_prob=0,
|
|
locations=True,
|
|
unblocking=False,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
class PickupDist(RoomGridLevel):
|
|
"""
|
|
|
|
## Description
|
|
|
|
Pick up an object
|
|
The object to pick up is given by its type only, or
|
|
by its color, or by its type and color.
|
|
(in the current room, with distractors)
|
|
|
|
## Mission Space
|
|
|
|
"pick up a/the {color}/{type}/{color}{type}"
|
|
|
|
{color} is the color of the box. Can be "red", "green", "blue", "purple",
|
|
"yellow" or "grey".
|
|
|
|
{type} is the type of the object. Can be "ball", "box" or "key".
|
|
|
|
## Action Space
|
|
|
|
| Num | Name | Action |
|
|
|-----|--------------|-------------------|
|
|
| 0 | left | Turn left |
|
|
| 1 | right | Turn right |
|
|
| 2 | forward | Move forward |
|
|
| 3 | pickup | Pick up an object |
|
|
| 4 | drop | Unused |
|
|
| 5 | toggle | Unused |
|
|
| 6 | done | Unused |
|
|
|
|
## Observation Encoding
|
|
|
|
- Each tile is encoded as a 3 dimensional tuple:
|
|
`(OBJECT_IDX, COLOR_IDX, STATE)`
|
|
- `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
|
|
[minigrid/minigrid.py](minigrid/minigrid.py)
|
|
- `STATE` refers to the door state with 0=open, 1=closed and 2=locked
|
|
|
|
## Rewards
|
|
|
|
A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
|
|
|
|
## Termination
|
|
|
|
The episode ends if any one of the following conditions is met:
|
|
|
|
1. The agent picks up the object.
|
|
2. Timeout (see `max_steps`).
|
|
|
|
## Registered Configurations
|
|
|
|
- `BabyAI-PickupDist-v0`
|
|
- `BabyAI-PickupDistDebug-v0`
|
|
|
|
"""
|
|
|
|
def __init__(self, debug=False, **kwargs):
|
|
self.debug = debug
|
|
super().__init__(num_rows=1, num_cols=1, room_size=7, **kwargs)
|
|
|
|
def gen_mission(self):
|
|
# Add 5 random objects in the room
|
|
objs = self.add_distractors(num_distractors=5)
|
|
self.place_agent(0, 0)
|
|
obj = self._rand_elem(objs)
|
|
type = obj.type
|
|
color = obj.color
|
|
|
|
select_by = self._rand_elem(["type", "color", "both"])
|
|
if select_by == "color":
|
|
type = None
|
|
elif select_by == "type":
|
|
color = None
|
|
|
|
self.instrs = PickupInstr(ObjDesc(type, color), strict=self.debug)
|
|
|
|
|
|
class PickupAbove(RoomGridLevel):
|
|
"""
|
|
|
|
## Description
|
|
|
|
Pick up an object (in the room above)
|
|
This task requires to use the compass to be solved effectively.
|
|
|
|
## Mission Space
|
|
|
|
"go to the {color} {type}"
|
|
|
|
{color} is the color of the box. Can be "red", "green", "blue", "purple",
|
|
"yellow" or "grey".
|
|
|
|
{type} is the type of the object. Can be "ball", "box" or "key".
|
|
|
|
## Action Space
|
|
|
|
| Num | Name | Action |
|
|
|-----|--------------|-------------------|
|
|
| 0 | left | Turn left |
|
|
| 1 | right | Turn right |
|
|
| 2 | forward | Move forward |
|
|
| 3 | pickup | Pick up an object |
|
|
| 4 | drop | Unused |
|
|
| 5 | toggle | Unused |
|
|
| 6 | done | Unused |
|
|
|
|
## Observation Encoding
|
|
|
|
- Each tile is encoded as a 3 dimensional tuple:
|
|
`(OBJECT_IDX, COLOR_IDX, STATE)`
|
|
- `OBJECT_TO_IDX` and `COLOR_TO_IDX` mapping can be found in
|
|
[minigrid/minigrid.py](minigrid/minigrid.py)
|
|
- `STATE` refers to the door state with 0=open, 1=closed and 2=locked
|
|
|
|
## Rewards
|
|
|
|
A reward of '1 - 0.9 * (step_count / max_steps)' is given for success, and '0' for failure.
|
|
|
|
## Termination
|
|
|
|
The episode ends if any one of the following conditions is met:
|
|
|
|
1. The agent picks up the object.
|
|
2. Timeout (see `max_steps`).
|
|
|
|
## Registered Configurations
|
|
|
|
- `BabyAI-PickupAbove-v0`
|
|
|
|
"""
|
|
|
|
def __init__(self, max_steps: int | None = None, **kwargs):
|
|
room_size = 6
|
|
if max_steps is None:
|
|
max_steps = 8 * room_size**2
|
|
|
|
super().__init__(room_size=room_size, max_steps=max_steps, **kwargs)
|
|
|
|
def gen_mission(self):
|
|
# Add a random object to the top-middle room
|
|
obj, pos = self.add_object(1, 0)
|
|
# Make sure the two rooms are directly connected
|
|
self.add_door(1, 1, 3, locked=False)
|
|
self.place_agent(1, 1)
|
|
self.connect_all()
|
|
|
|
self.instrs = PickupInstr(ObjDesc(obj.type, obj.color))
|