Package: memory-maze
Status:Stable releaseMemory MazeMemory 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 9x9Memory 11x11Memory 13x13Memory 15x15Key features:Online RL memory tasks (with baselines)Offline dataset for representation learning (with baselines)Verified that memory is the key challengeChallenging but solvable by human baselineEasy installation via a simple pip commandAvailablegymanddm_envinterfacesSupports headless and hardware renderingInteractive GUI for human playersHidden state information for probe evaluationAlso 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}
}InstallationMemory Maze builds on thedm_controlandmujocopackages, which are automatically installed as dependencies:pipinstallmemory-mazePlay YourselfMemory 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 dependenciespipinstallgympygamepillowimageio# Launch with standard 64x64 resolutionpythongui/run_gui.py# Launch with higher 256x256 resolutionpythongui/run_gui.py--env"memory_maze:MemoryMaze-9x9-HD-v0"Task DescriptionThe 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.Sizeenv_idObjectsEpisode stepsMean human scoreMean max score9x9MemoryMaze-9x9-v03100026.434.811x11MemoryMaze-11x11-v04200044.358.013x13MemoryMaze-13x13-v05300055.574.515x15MemoryMaze-15x15-v06400067.787.7The mazes are generated withlabmaze, the same algorithm as used byDmLab-30. The 9x9 corresponds to thesmallvariant and 15x15 corresponds to thelargevariant.Gym InterfaceYou can create the environment using theGyminterface:!pipinstallgymimportgym# 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 hereabout the different rendering options.The default environment has 64x64 image observations:>>>env.observation_spaceBox(0,255,(64,64,3),uint8)There are 6 discrete actions:>>>env.action_spaceDiscrete(6)# (noop, forward, left, right, forward_left, forward_right)To create an environment with extra observations for debugging and probe analysis, appendExtraObsto the names:>>>env=gym.make('memory_maze:MemoryMaze-9x9-ExtraObs-v0')>>>env.observation_spaceDict(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 registeradditional variantsof the environment that can be useful in certain scenarios.DeepMind InterfaceYou can create the environment using thedm_envinterface:frommemory_mazeimporttasksenv=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 includesimagekey:>>>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,)Offline DatasetDataset 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:KeyShapeTypeDescriptionimage(64, 64, 3)uint8First-person view observationaction(6)binaryLast action, one-hot encodedreward()floatLast rewardmaze_layout(9, 9) or (15, 15)binaryMaze layout (wall / no wall)agent_pos(2)floatAgent position in global coordinatesagent_dir(2)floatAgent orientation as a unit vectortargets_pos(3, 2) or (6, 2)floatObject locations in global coordinatestargets_vec(3, 2) or (6, 2)floatObject locations in agent-centric coordinatestarget_pos(2)floatCurrent target object location, globaltarget_vec(2)floatCurrent target object location, agent-centrictarget_color(3)floatCurrent target object color RGBYou can load a trajectory usingnp.load()to obtain a dictionary of Numpy arrays as follows:episode=np.load('trajectory.npz')episode={key:episode[key]forkeyinepisode.keys()}assertepisode['image'].shape==(1001,64,64,3)assertepisode['image'].dtype==np.uint8All tensors have a leading time dimension, e.g.imagetensor has shape (1001, 64, 64, 3). The tensor length is 1001 because there are 1000 steps (actions) in a trajectory,image[0]is the observationbeforethe first action, andimage[-1]is the observationafterthe last action.Online RL BaselinesIn ourresearch paper, we evaluate the model-freeIMPALAagent and the model-basedDreameragent as baselines.Here are videos of the learned behaviors:Memory 9x9 - Dreamer (TBTT)https://user-images.githubusercontent.com/3135115/197378287-4e413440-7097-4d11-8627-3d7fac0845f1.mp4Memory 9x9 - IMPALA (400M)https://user-images.githubusercontent.com/3135115/197378929-7fe3f374-c11c-409a-8a95-03feeb489330.mp4Memory 15x15 - Dreamer (TBTT)https://user-images.githubusercontent.com/3135115/197378324-fb99b496-dba8-4b00-ad80-2d6e19ba8acd.mp4Memory 15x15 - IMPALA (400M)https://user-images.githubusercontent.com/3135115/197378936-939e7615-9dad-4765-b0ef-a49c5a38fe28.mp4Offline Probing BaselinesHere we visualize probe predictions alongside trajectories of the offline dataset, as explained inthe paper. These trajectories are from the offline dataset, where the agent just navigates to random points in the maze, it doesnottry 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)https://user-images.githubusercontent.com/3135115/197379227-775ec5bc-0780-4dcc-b7f1-660bc7cf95f1.mp4Memory 9x9 Walls Objects - Supervised oraclehttps://user-images.githubusercontent.com/3135115/197379235-a5ea0388-2718-4035-8bbc-064ecc9ea444.mp4Memory 15x15 Walls Objects - RSSM (TBTT)https://user-images.githubusercontent.com/3135115/197379245-fb96bd12-6ef5-481e-adc6-f119a39e8e43.mp4Memory 15x15 Walls Objects - Supervised oraclehttps://user-images.githubusercontent.com/3135115/197379248-26a8093e-8b54-443c-b154-e33e0383b5e4.mp4QuestionsPleaseopen an issueon Github.