In-Context Reinforcement Learning From Suboptimal Historical Data

TL;DR

This paper introduces the Decision Importance Transformer (DIT), a novel framework that enables in-context reinforcement learning from suboptimal historical data by emulating actor-critic algorithms with transformers.

Contribution

The paper proposes DIT, a transformer-based method that estimates advantage functions and improves policy learning from suboptimal offline datasets in RL tasks.

Findings

DIT outperforms standard methods on bandit and MDP problems.

The approach effectively leverages suboptimal data to learn near-optimal policies.

Transformers can emulate actor-critic algorithms in an in-context learning setting.

Abstract

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.