Unitary recurrent networks improve memory in reinforcement learning tasks
Reinforcement Learning with Complex (valued) Memories
Machine Learning
Summary
Remembering past information to make decisions is hard for many AI systems when they can't see everything at once. The authors looked back at a special kind of AI method called unitary recurrent networks that store information in a way that doesn't lose details over time. They adapted these networks to improve how AI learns tasks where memory is important, like controlling robots or playing certain games. Their methods performed much better on several tests, showing that using complex numbers to represent memory helps AI remember longer and make better decisions.
What this means in practice
- •For robotics engineers: Improve robot decision-making in tasks where sensors provide incomplete or indirect observations by using unitary recurrent networks for better long-term memory.
- •For game ai developers: Enhance AI agents in games involving hidden information or delayed rewards by integrating complex-valued memory structures to maintain critical past observations.
Authors
Sathya Kamesh Bhethanabhotla, Efstratios Gavves, André Biedenkapp
Abstract
Partially observable environments pose a fundamental challenge in deep reinforcement learning, requiring agents to compress temporal information from observations and maintain a memory to make effective decisions. While there exist many approaches ranging from gated recurrence to attention mechanisms and model-based RL, the search for effective representational techniques that can capture long-term dependencies remains an active area of research. In this work we revisit Unitary recurrent networks (uRNNs) [Arjovsky et al., 2016, Jing et al., 2017], that demonstrated superior gradient flow and associative recall, expressing the recurrence and the hidden state in a complex vector space. Their norm preserving unitary dynamics enable information propagation through long sequences. To this end, we propose three different versions of uRNNs as drop-in replacements for recurrent PPO architectures, and demonstrate that the simple recurrence and the added degree of freedom from the phase of the complex representations enable significant gains over baselines on several memory-improvable tasks, including continuous control. We further explore how to preserve the phase information of the complex hidden state for a phase-aware policy by drawing a parallel to how quantum states are measured. With our methods reaching up to 2-3 $\times$ the reward in environments like rocksample and Craftax compared to the baselines, this work points towards an exciting new direction of representations for RL and the problem of partial observability. Code is available at: https://github.com/Sathya98/qurl