Status: Stable release
Memory Maze is a 3D domain of randomized mazes designed for evaluating the long-term memory abilities of RL agents. Memory Maze isolates long-term memory from confounding challenges, such as exploration, and requires remembering several pieces of information: the positions of objects, the wall layout, and keeping track of agent’s own position.
| Memory 9x9 | Memory 11x11 | Memory 13x13 | Memory 15x15 |
|---|---|---|---|
 |
 |
 |
 |
Key features:
- Online RL memory tasks (with baselines)
- Offline dataset for representation learning (with baselines)
- Verified that memory is the key challenge
- Challenging but solvable by human baseline
- Easy installation via a simple pip command
- Available
gymanddm_envinterfaces - Supports headless and hardware rendering
- Interactive GUI for human players
- Hidden state information for probe evaluation
Also see the accompanying research paper: Evaluating Long-Term Memory in 3D Mazes
@article{pasukonis2022memmaze,
title={Evaluating Long-Term Memory in 3D Mazes},
author={Pasukonis, Jurgis and Lillicrap, Timothy and Hafner, Danijar},
journal={arXiv preprint arXiv:2210.13383},
year={2022}
}
Memory Maze builds on the dm_control and mujoco packages, which are automatically installed as dependencies:
pip install memory-mazeMemory Maze allows you to play the levels in human mode. We used this mode for recording the human baseline scores. These are the instructions for launching the GUI:
# GUI dependencies
pip install gym pygame pillow imageio
# Launch with standard 64x64 resolution
python gui/run_gui.py
# Launch with higher 256x256 resolution
python gui/run_gui.py --env "memory_maze:MemoryMaze-9x9-HD-v0"The task is based on a game known as scavenger hunt or treasure hunt:
- The agent starts in a randomly generated maze, which contains several objects of different colors.
- The agent is prompted to find the target object of a specific color, indicated by the border color in the observation image.
- Once the agent successfully finds and touches the correct object, it gets a +1 reward and the next random object is chosen as a target.
- If the agent touches an object of the wrong color, there is no effect.
- Throughout the episode, the maze layout and the locations of the objects do not change.
- The episode continues for a fixed amount of time, so the total episode reward equals the number of reached targets.
An agent with long-term memory only has to explore each maze once (which is possible in a time much shorter than the length of an episode) and can afterwards follow the shortest path to each requested target, whereas an agent with no memory has to randomly wander through the maze to find each target.
There are 4 size variations of the maze. The largest maze 15x15 is designed to be challenging but solvable for humans (see benchmark results below), but out of reach for the state-of-the-art RL methods. The smaller sizes are provided as stepping stones, with 9x9 being solvable with current RL methods.
| Size | env_id | Objects | Episode steps | Mean human score | Mean max score |
|---|---|---|---|---|---|
| 9x9 | MemoryMaze-9x9-v0 |
3 | 1000 | 26.4 | 34.8 |
| 11x11 | MemoryMaze-11x11-v0 |
4 | 2000 | 44.3 | 58.0 |
| 13x13 | MemoryMaze-13x13-v0 |
5 | 3000 | 55.5 | 74.5 |
| 15x15 | MemoryMaze-15x15-v0 |
6 | 4000 | 67.7 | 87.7 |
The mazes are generated with labmaze, the same algorithm as used by DmLab-30. The 9x9 corresponds to the small variant and 15x15 corresponds to the large variant.
You can create the environment using the Gym interface:
!pip install gym
import gym
# Set this if you are getting "Unable to load EGL library" error:
# os.environ['MUJOCO_GL'] = 'glfw'
env = gym.make('memory_maze:MemoryMaze-9x9-v0')
env = gym.make('memory_maze:MemoryMaze-11x11-v0')
env = gym.make('memory_maze:MemoryMaze-13x13-v0')
env = gym.make('memory_maze:MemoryMaze-15x15-v0')Troubleshooting: if you are getting "Unable to load EGL library error", that is because we enable MuJoCo headless GPU rendering (MUJOCO_GL=egl) by default. If you are testing locally on your machine, you can enable windowed rendering instead (MUJOCO_GL=glfw). Read here about the different rendering options.
The default environment has 64x64 image observations:
>>> env.observation_space
Box(0, 255, (64, 64, 3), uint8)There are 6 discrete actions:
>>> env.action_space
Discrete(6) # (noop, forward, left, right, forward_left, forward_right)To create an environment with extra observations for debugging and probe analysis, append ExtraObs to the names:
>>> env = gym.make('memory_maze:MemoryMaze-9x9-ExtraObs-v0')
>>> env.observation_space
Dict(
agent_dir: Box(-inf, inf, (2,), float64),
agent_pos: Box(-inf, inf, (2,), float64),
image: Box(0, 255, (64, 64, 3), uint8),
maze_layout: Box(0, 1, (9, 9), uint8),
target_color: Box(-inf, inf, (3,), float64),
target_pos: Box(-inf, inf, (2,), float64),
target_vec: Box(-inf, inf, (2,), float64),
targets_pos: Box(-inf, inf, (3, 2), float64),
targets_vec: Box(-inf, inf, (3, 2), float64)
)We also register additional variants of the environment that can be useful in certain scenarios.
You can create the environment using the dm_env interface:
from memory_maze import tasks
env = tasks.memory_maze_9x9()
env = tasks.memory_maze_11x11()
env = tasks.memory_maze_13x13()
env = tasks.memory_maze_15x15()Each observation is a dictionary that includes image key:
>>> env.observation_spec()
{
'image': BoundedArray(shape=(64, 64, 3), ...)
}The constructor accepts a number of arguments, which can be used to tweak the environment:
env = tasks.memory_maze_9x9(
global_observables=True,
image_only_obs=False,
top_camera=False,
camera_resolution=64,
control_freq=4.0,
discrete_actions=True,
)Dataset download here (~100GB per dataset)
We provide two datasets of experience collected from the Memory Maze environment: Memory Maze 9x9 (30M) and Memory Maze 15x15 (30M). Each dataset contains 30 thousand trajectories from Memory Maze 9x9 and 15x15 environments respectively, split into 29k trajectories for training and 1k for evaluation. All trajectories are 1000 steps long, so each dataset has 30M steps total.
The data is generated with a scripted policy that navigates to randomly chosen points in the maze under action noise. This choice of policy was made to generate diverse trajectories that explore the maze effectively and that form spatial loops, which can be important for learning long-term memory. We intentionally avoid recording data with a trained agent to ensure a diverse data distribution and to avoid dataset bias that could favor some methods over others. Because of this, the rewards are quite sparse in the data, occurring on average 1-2 times per trajectory.
Each trajectory is saved as an NPZ file with the following entries available:
| Key | Shape | Type | Description |
|---|---|---|---|
image |
(64, 64, 3) | uint8 | First-person view observation |
action |
(6) | binary | Last action, one-hot encoded |
reward |
() | float | Last reward |
maze_layout |
(9, 9) or (15, 15) | binary | Maze layout (wall / no wall) |
agent_pos |
(2) | float | Agent position in global coordinates |
agent_dir |
(2) | float | Agent orientation as a unit vector |
targets_pos |
(3, 2) or (6, 2) | float | Object locations in global coordinates |
targets_vec |
(3, 2) or (6, 2) | float | Object locations in agent-centric coordinates |
target_pos |
(2) | float | Current target object location, global |
target_vec |
(2) | float | Current target object location, agent-centric |
target_color |
(3) | float | Current target object color RGB |
You can load a trajectory using np.load() to obtain a dictionary of Numpy arrays as follows:
episode = np.load('trajectory.npz')
episode = {key: episode[key] for key in episode.keys()}
assert episode['image'].shape == (1001, 64, 64, 3)
assert episode['image'].dtype == np.uint8All tensors have a leading time dimension, e.g. image tensor has shape (1001, 64, 64, 3). The tensor length is 1001 because there are 1000 steps (actions) in a trajectory, image[0] is the observation before the first action, and image[-1] is the observation after the last action.
In our research paper, we evaluate the model-free IMPALA agent and the model-based Dreamer agent as baselines.
Here are videos of the learned behaviors:
Memory 9x9 - Dreamer (TBTT)
rl-9x9-dreamer-r32.mp4
Memory 9x9 - IMPALA (400M)
rl-9x9-impala-r23.mp4
Memory 15x15 - Dreamer (TBTT)
rl-15x15-dreamer-r1-1000.mp4
Memory 15x15 - IMPALA (400M)
rl-15x15-impala-r25.mp4
Here we visualize probe predictions alongside trajectories of the offline dataset, as explained in the paper. These trajectories are from the offline dataset, where the agent just navigates to random points in the maze, it does not try to collect rewards.
Bottom-left: Object location predictions (x) versus the actual locations (o).
Bottom-right: Wall layout predictions (dark green = true positive, light green = true negative, light red = false positive, dark red = false negative).
Memory 9x9 Walls Objects - RSSM (TBTT)
probing-9x9-dreamer.mp4
Memory 9x9 Walls Objects - Supervised oracle
probing-9x9-oracle.mp4
Memory 15x15 Walls Objects - RSSM (TBTT)
probing-15x15-dreamer.mp4
Memory 15x15 Walls Objects - Supervised oracle
probing-15x15-oracle.mp4
Please open an issue on Github.