In-Context Reinforcement Learning From Suboptimal Historical Data
This addresses the challenge of leveraging suboptimal offline data for reinforcement learning, offering a method to enhance policy learning in new tasks, though it is incremental as it builds on existing transformer and actor-critic approaches.
The paper tackles the problem of in-context reinforcement learning from suboptimal historical data by proposing the Decision Importance Transformer (DIT) framework, which emulates actor-critic algorithms to improve policy performance, achieving superior results in experiments on bandit and Markov Decision Process problems.
Transformer models have achieved remarkable empirical successes, largely due to their in-context learning capabilities. Inspired by this, we explore training an autoregressive transformer for in-context reinforcement learning (ICRL). In this setting, we initially train a transformer on an offline dataset consisting of trajectories collected from various RL tasks, and then fix and use this transformer to create an action policy for new RL tasks. Notably, we consider the setting where the offline dataset contains trajectories sampled from suboptimal behavioral policies. In this case, standard autoregressive training corresponds to imitation learning and results in suboptimal performance. To address this, we propose the Decision Importance Transformer(DIT) framework, which emulates the actor-critic algorithm in an in-context manner. In particular, we first train a transformer-based value function that estimates the advantage functions of the behavior policies that collected the suboptimal trajectories. Then we train a transformer-based policy via a weighted maximum likelihood estimation loss, where the weights are constructed based on the trained value function to steer the suboptimal policies to the optimal ones. We conduct extensive experiments to test the performance of DIT on both bandit and Markov Decision Process problems. Our results show that DIT achieves superior performance, particularly when the offline dataset contains suboptimal historical data.