Decision Titan improves long-term memory in offline reinforcement learning
Decision Titan: Test-Time Training for Long-Term Memory in Offline Reinforcement Learning
Artificial IntelligenceMachine Learning
Summary
Remembering information over long times is hard for AI systems that make decisions step-by-step. The authors studied a new idea called Test-Time Training that helps a kind of AI called Decision Transformer keep better memory by adjusting itself during use. They tested this approach in a maze environment designed to challenge memory and found it could remember things 20 times longer than usual. The improvements depend on how time and information are represented inside the system.
What this means in practice
- •For robotics engineers: Enhance robots’ decision-making by enabling memory over longer time spans in offline learning scenarios using the Decision Titan approach.
- •For game ai developers: Improve AI agents’ ability to remember past actions and plan better in complex sequential games by integrating test-time training with decision transformers.
Tested on simulated data.
Authors
Jude Waide, Robert Lieck
Abstract
Long-term dependencies remain a major challenge for sequential decision-making in the field of AI: RNNs suffer from vanishing gradients and the limited expressivity of vector-based hidden states, whilst Transformer-based models are limited by the quadratic scaling of attention. Recent work has proposed tackling this problem with the Test-Time Training (TTT) framework, which stores episodic memories in the parameters of a neural network through gradient descent at both train and test-time. This approach has seen success in the domain of Natural Language Processing, however, to the best of our knowledge it has not yet been applied to the domain of Reinforcement Learning (RL), nor has there been a study analysing how this memory practically functions. In this paper, we study the potential of the TTT framework for offline RL by augmenting a Decision Transformer with TTT layers, dubbed the Decision Titan. We analyse performance and properties of the model in the X-Maze environment, an extension of T-Maze designed to test sequential memory, and investigate how the memory mechanism learns by visualising gate values over time. Our key findings are that Decision Titan can learn long-term dependencies with ranges 20x longer than the context window, generalises to lengths 1.7x the training data, but crucially temporal generalisation depends on the time embeddings used, and the ability to learn long-term dependencies depends on how the relevant information is encoded.