gym-balletenv

The gym-balletenv is a Gym-style implementation of the Ballet environment introduced in the paper “Towards mental time travel: a hierarchical memory for reinforcement learning agents,” designed to test the HCAM architecture. This environment allows users to experiment with memory architectures, offering a robust platform to test agents’ ability to handle complex tasks that require long-term memory and decision-making. The Hierarchical Chunk Attention Memory (HCAM) is a novel memory architecture proposed to address the limitations of traditional reinforcement learning (RL) agents, which often struggle with long-term recall and integration of events over multiple timesteps or following distractor tasks. HCAM allows agents to “mentally time-travel” by dividing past events into chunks and recalling them through high-level attention over coarse summaries, followed by detailed attention on the most relevant chunks. This enables agents to retain and reason over long-term memory with greater efficiency and generalization, excelling at tasks like recalling object locations in 3D environments, navigating new areas, and learning object names. HCAM significantly improves the sample efficiency, generalization, and adaptability of RL agents, pushing the boundaries of memory-based learning in temporally-extended environments.