Gym-MiniGrid/gym_minigrid/envs/redbluedoors.py

from gym_minigrid.minigrid import Door, Grid, MiniGridEnv
from gym_minigrid.register import register


class RedBlueDoorEnv(MiniGridEnv):
    """
    Single room with red and blue doors on opposite sides.
    The red door must be opened before the blue door to
    obtain a reward.
    """

    def __init__(self, size=8):
        self.size = size

        super().__init__(width=2 * size, height=size, max_steps=20 * size * size)

    def _gen_grid(self, width, height):
        # Create an empty grid
        self.grid = Grid(width, height)

        # Generate the grid walls
        self.grid.wall_rect(0, 0, 2 * self.size, self.size)
        self.grid.wall_rect(self.size // 2, 0, self.size, self.size)

        # Place the agent in the top-left corner
        self.place_agent(top=(self.size // 2, 0), size=(self.size, self.size))

        # Add a red door at a random position in the left wall
        pos = self._rand_int(1, self.size - 1)
        self.red_door = Door("red")
        self.grid.set(self.size // 2, pos, self.red_door)

        # Add a blue door at a random position in the right wall
        pos = self._rand_int(1, self.size - 1)
        self.blue_door = Door("blue")
        self.grid.set(self.size // 2 + self.size - 1, pos, self.blue_door)

        # Generate the mission string
        self.mission = "open the red door then the blue door"

    def step(self, action):
        red_door_opened_before = self.red_door.is_open
        blue_door_opened_before = self.blue_door.is_open

        obs, reward, done, info = MiniGridEnv.step(self, action)

        red_door_opened_after = self.red_door.is_open
        blue_door_opened_after = self.blue_door.is_open

        if blue_door_opened_after:
            if red_door_opened_before:
                reward = self._reward()
                done = True
            else:
                reward = 0
                done = True

        elif red_door_opened_after:
            if blue_door_opened_before:
                reward = 0
                done = True

        return obs, reward, done, info


class RedBlueDoorEnv6x6(RedBlueDoorEnv):
    def __init__(self):
        super().__init__(size=6)


register(
    id="MiniGrid-RedBlueDoors-6x6-v0", entry_point="gym_minigrid.envs:RedBlueDoorEnv6x6"
)

register(
    id="MiniGrid-RedBlueDoors-8x8-v0", entry_point="gym_minigrid.envs:RedBlueDoorEnv"
)