Complex Skill Learning using Modified Transformers with Automatic Chunking and Forgetting

Transformers are increasingly becoming mainstream in the field of reinforcement learning due to their ability to learn complex behaviour sequences over long periods of time. Human learning and decision-making works on recollection which is based on only parts of memory that are important to the current context. The past contains information that might not be useful for making decisions in the present. The human brain reduces its need to store every past event by forgetting less important events and only associating importance to the events that lead to the desired outcome, giving a broader picture of the past and helping with faster decision making. Current transformer models lack the ability to scale to large memory sizes and learn long horizon tasks efficiently as the computational requirements needed to run these models scale non-linearly with memory length. We propose Automatic Chunking + ForgetSpan, an architecture that reduces the computational performance needed to train and run transformer models by leveraging automatic chunking and forgetting. Automatic chunking helps with chunking the memory into parts which might be relevant in the present context while forgetting removes events within the memory that do not contribute to learning, helping the model learn multiple sequences of events with ease and reducing the computational cost of larger memory sizes. We demonstrate that Automatic Chunking + ForgetSpan can help models to memorize important information and achieve improved performance with reduced computational needs on various memory, visual navigation, robotic locomotion and multimodal tasks. We also test SimilarityWeight, a method which weighs a new memory element based on its similarity with already existing elements in the buffer. Comparative analysis with existing memory architectures was performed to understand the effects of relevant memory selection and forgetting. Gating connections which make use of gated recurrent units and memory mean pooling based on masking are also used for further improvement in performance. Effects of using multiple memory-based transformers with automatic chunking in series operating on data from consecutive timesteps were also studied.