"""The Simple Maze MDP environment (`MazeMDP-v0`).
A small tabular Markov Decision Process on a grid, used by the
dynamic-programming and tabular-RL notebooks. Wraps the ``mazemdp`` package
and exposes the transition model (``P``, ``r``, ``terminal_states``, ``gamma``)
for planning, plus gymnasium-style ``reset`` / ``step`` for learning.
Importing this module registers ``MazeMDP-v0`` in gymnasium. Adapted from
``bbrl_gymnasium`` (MIT license, Sorbonne Université).
"""
from __future__ import annotations
import logging
from functools import partial
from typing import Any, Callable
import gymnasium as gym
from gymnasium import spaces
from gymnasium.core import ObsType
from gymnasium.utils import seeding
from mazemdp import create_random_maze
from mazemdp.maze import build_custom_maze
logger = logging.getLogger(__name__)
__all__ = ["MazeMDPEnv"]
[docs]
class MazeMDPEnv(gym.Env):
"""A tabular grid-world MDP, with its transition model exposed
Observations are **state indices** (a `Discrete` space) and there are four
actions (north, south, east, west). Unlike a usual gymnasium environment,
the whole model is available as attributes (`P`, `r`, `terminal_states`,
`gamma`), which is what the dynamic-programming notebooks need; the
`draw_*` methods render a value function and/or a policy on the grid.
Created through gymnasium, either random or from an explicit wall list::
import rl_mind.mazemdp_env # registers MazeMDP-v0
env = gym.make("MazeMDP-v0", kwargs={"width": 5, "height": 5,
"ratio": 0.2, "seed": 0})
:param render_mode: `"rgb_array"`, `"human"`, or `None`
:param kwargs: A single `kwargs=` dictionary describing the maze:
`width`, `height`, `hit` (reward for hitting a wall), then either
`ratio` (fraction of walls, with an optional `seed` for a reproducible
random maze) or `walls` (explicit wall indices, with optional
`terminal_states`), plus optional `start_states`. With no argument at
all, a random 10x10 maze with 20% walls is built.
"""
metadata = {"render_modes": ["rgb_array", "human"], "video.frames_per_second": 5}
def __init__(self, render_mode=None, **kwargs):
self.render_mode = render_mode
if kwargs == {}:
width = height = 10
self.mdp, nb_states, coord_x, coord_y = create_random_maze(10, 10, 0.2, [0])
else:
kwargs = kwargs["kwargs"]
width = kwargs["width"]
height = kwargs["height"]
hit = kwargs.get("hit", False)
if "walls" not in kwargs:
ratio = kwargs["ratio"]
walls = None
else:
ratio = None
walls = kwargs["walls"]
terminal_states = kwargs.get("terminal_states", None)
if terminal_states is not None and walls is None:
print(
"warning: one should not set terminal states in a random maze "
"as the final state might contain a wall"
)
start_states = kwargs.get("start_states", None)
seed = kwargs.get("seed", None)
if walls is None:
# A seeded random maze is drawn from a private generator, so
# that every construction with the same seed gives the *same*
# maze -- which is what `VecEnv` needs, since it builds one
# copy of the environment per environment (requires mazemdp
# >= 1.3.0; without a seed, any version works)
seeding = {} if seed is None else {"seed": seed}
self.mdp, nb_states, coord_x, coord_y = create_random_maze(
width,
height,
hit,
ratio=ratio,
start_states=start_states,
**seeding,
)
else:
self.mdp, nb_states, coord_x, coord_y = build_custom_maze(
width,
height,
hit,
walls=walls,
terminal_states=terminal_states,
start_states=start_states,
)
#: The number of states (walls excluded)
self.nb_states = nb_states
#: The grid column of each state
self.coord_x = coord_x
#: The grid row of each state
self.coord_y = coord_y
#: The observation space: a state index in `[0, nb_states)`
self.observation_space = spaces.Discrete(nb_states)
#: The action space: 0 = north, 1 = south, 2 = east, 3 = west
self.action_space = spaces.Discrete(4)
#: The indices of the terminal states
self.terminal_states = self.mdp.terminal_states
#: The transition model $P(s' \mid s, a)$, shape
#: `[nb_states, 4, nb_states]`
self.P = self.mdp.P
#: The discount factor $\gamma$ of the MDP
self.gamma = self.mdp.gamma
#: The reward model $r(s, a)$, shape `[nb_states, 4]`
self.r = self.mdp.r
#: The current state index. Reading it is cheating: a learning agent
#: only knows the states `reset` / `step` return.
self.state = self.mdp.current_state
self.seed()
self.np_random = None
self.title = f"Simple maze {width}x{height}"
self.set_render_func(partial(self.mdp.render, title=self.title))
[docs]
def set_title(self, title):
"""Set the title drawn above the maze
:param title: The new title
"""
self.title = title
[docs]
def set_render_func(self, render_func: Callable, *args, **kwargs):
"""Set the function `render` calls
:param render_func: The drawing function (the `mazemdp` renderer by
default); the current `render_mode` is passed as `mode=`
:param args: Positional arguments bound to `render_func`
:param kwargs: Keyword arguments bound to `render_func`
"""
kwargs = {**kwargs, "mode": self.render_mode}
self.render_func = partial(render_func, *args, **kwargs)
[docs]
def seed(self, seed=None):
"""Seed the environment's random generator
:param seed: The seed, or `None` to draw one
:returns: A one-element list holding the seed actually used
"""
self.np_random, seed = seeding.np_random(seed)
return [seed]
[docs]
def step(self, action):
"""Perform one action in the maze
:param action: The action index (0 = north, 1 = south, 2 = east,
3 = west)
:returns: The gymnasium 5-tuple `(state, reward, terminated,
truncated, info)`, where `state` is a state index
"""
return self.mdp.step(action)
[docs]
def reset(
self, *, seed: int | None = None, options: dict[str, Any] | None = None
) -> tuple[ObsType, dict[str, Any]]:
"""Start a new episode
:param seed: Unused (the maze itself is fixed at construction time;
seed it there for a reproducible random maze)
:param options: Passed to `mazemdp`'s `reset`, e.g.
`{"uniform": True}` to start from a state drawn uniformly
:returns: The initial state index and an (empty) info dictionary
"""
options = options or {}
r = self.mdp.reset(**options)
if isinstance(r, list):
return r
return self.mdp.reset(**options), {}
[docs]
def sample_transition(self, **kwargs):
"""Draw a transition from the model, without stepping the environment
:param kwargs: Passed to `mazemdp`'s `sample_transition`
:returns: What `mazemdp` returns, a sampled `(state, action, reward,
next state)` transition
"""
return self.mdp.sample_transition(**kwargs)
def _draw(self, recorder, callable, *args, **kwargs):
"""Draw and optionally record a frame"""
if recorder is not None:
if recorder.enabled:
self.set_render_func(callable, *args, **kwargs)
recorder.capture_frame()
else:
return callable(*args, **kwargs)
[docs]
def draw_v_pi_a(self, v, policy, agent_pos, title="MDP studies", recorder=None):
"""Draw the value function, the policy and the agent's position
:param v: The value of each state, shape `[nb_states]`
:param policy: The action chosen in each state, shape `[nb_states]`
:param agent_pos: The state the agent stands on
:param title: The title drawn above the maze
:param recorder: A video recorder to capture the frame into, instead
of displaying it
:returns: What the renderer returns (an `[H, W, 3]` array in
`rgb_array` mode), or `None` when recording
"""
return self._draw(
recorder,
self.render,
v,
policy,
agent_pos,
title=title,
mode=self.render_mode,
)
[docs]
def draw_v_pi(self, v, policy, title="MDP studies", recorder=None):
"""Draw the value function (colors) and the policy (arrows)
:param v: The value of each state, shape `[nb_states]`
:param policy: The action chosen in each state, shape `[nb_states]`
:param title: The title drawn above the maze
:param recorder: A video recorder to capture the frame into, instead
of displaying it
:returns: What the renderer returns, or `None` when recording
"""
return self._draw(
recorder, self.mdp.render, v, policy, None, title, mode=self.render_mode
)
[docs]
def draw_v(self, v, title="MDP studies", recorder=None):
"""Draw a value function (or the Q-values of each action)
:param v: The value of each state (shape `[nb_states]`), or of each
state-action pair (shape `[nb_states, 4]`, drawn as four
triangles per cell)
:param title: The title drawn above the maze
:param recorder: A video recorder to capture the frame into, instead
of displaying it
:returns: What the renderer returns, or `None` when recording
"""
return self._draw(
recorder, self.mdp.render, v, None, None, title, mode=self.render_mode
)
[docs]
def draw_pi(self, policy, title="MDP studies", recorder=None):
"""Draw a policy (one arrow per state)
:param policy: The action chosen in each state, shape `[nb_states]`
:param title: The title drawn above the maze
:param recorder: A video recorder to capture the frame into, instead
of displaying it
:returns: What the renderer returns, or `None` when recording
"""
return self._draw(
recorder, self.mdp.render, None, policy, None, title, mode=self.render_mode
)
[docs]
def init_draw(self, title, recorder=None):
"""Open (or reset) the figure the maze is drawn in
:param title: The title drawn above the maze
:param recorder: A video recorder to capture the frame into
:returns: What the renderer returns, or `None` when recording
"""
return self._draw(recorder, self.mdp.new_render, title, mode=self.render_mode)
[docs]
def render(self):
"""Render the maze with the current render function
:returns: An `[H, W, 3]` array in `rgb_array` mode, `None` in
`human` mode
"""
return self.render_func(mode=self.render_mode)
[docs]
def render_mdp(self, func, *args, **kwargs):
"""Render with a custom drawing function (legacy; prefer
`set_render_func` followed by `render`)
Note that this *replaces* the `set_render_func` method with a partial
of `func`, so later calls to it fail; nothing in the practicals uses
this method.
:param func: The drawing function
:param args: Positional arguments bound to `func`
:param kwargs: Keyword arguments bound to `func`
"""
self.set_render_func = partial(func, *args, **kwargs)
self.render()
[docs]
def set_no_agent(self):
"""Stop drawing the agent (for pure dynamic-programming figures)"""
self.mdp.has_state = False
[docs]
def set_timeout(self, timeout):
"""Set the maximum number of steps of an episode
:param timeout: The step limit after which the episode is truncated
"""
self.mdp.timeout = timeout
# Importing this module registers the environment (explicit opt-in)
if "MazeMDP-v0" not in gym.registry:
gym.register(id="MazeMDP-v0", entry_point=MazeMDPEnv, max_episode_steps=200)