RI: SMALL: Improved Learning in Partially Observable Reinforcement Learning with Attention-Based Models

Reinforcement learning (RL), where machines learn to perform tasks by trial and error, has become an incredibly important paradigm in artificial intelligence. Nevertheless, much of the progress has focused on fully observable domains, where we have full knowledge of the world that these machines are in. Real-world domains are typically partially observable with limited information as to what is happening.

For example, domains such as autonomous driving or robotics are partially observable due to limited ability to sense what is happening in the environment. In order for reinforcement learning to perform well in these partially observable applications, machines must be able to remember relevant information from the past. Current methods typically use relatively simple approaches based on neural networks to determine what information to remember and what to forget.

In contrast, algorithms, like transformers, have shown a massive improvement over other models in language (e.g., ChatGPT) and vision (e.g., DALL-E) problems. This project will develop new methods for partially observable information, which should lead to significantly improved performance in partially observable reinforcement learning problems.

This project will develop a number of novel methods for partial observable reinforcement learning (PORL). In particular, the project will develop attention-based model-free approaches (including the first attention-based actor-critic methods) that perform well, are computationally efficient, can reason over long horizons, incorporate structure such as permutation invariant histories, and can process high-dimensional vision input.

The project will also develop the first attention-based model-based PORL methods. These approaches will improve sample efficiency and performance by incorporating supervised auxiliary tasks, dynamics models, or additional planning steps for improving exploration and model learning. Lastly, the project will formalize goal-conditioning in partially observable Markov decision processes (POMDPs), develop the first attention-based multi-task and goal-conditioned PORL methods, and develop pre-training approaches that can train on a wide range of domains and fine-tune to desired test domains.

The resulting methods and analysis will allow the community to better understand the role of attention-based methods for PORL and to build on this work to develop approaches that can efficiently learn in realistic partially observable domains.

This award reflects NSF's statutory mission and has been deemed worthy of support through evaluation using the Foundation's intellectual merit and broader impacts review criteria.