Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

In-Context Reinforcement Learning From Suboptimal Historical Data

Juncheng Dong, Moyang Guo, Ethan X. Fang, Zhuoran Yang, Vahid Tarokh

TL;DR

This work tackles in-context reinforcement learning (ICRL) from suboptimal historical data by introducing the Decision Importance Transformer (DIT). DIT pretrains a task-conditioned transformer on trajectories from diverse RL tasks and uses in-context advantage estimation to weight actions during a weighted maximum likelihood objective, enabling policy improvement without requiring optimal action labels. A key contribution is the in-context estimator setup—two transformers modelled to output in-context $V$ and $Q$, producing $\widehat{A}_{b}(s,a| au)$ used to form exponential weights for pretraining. Empirically, DIT matches or outperforms related approaches on both bandit and multi-task MDp benchmarks, demonstrating strong generalization to unseen tasks and robustness to suboptimal data, thus broadening the practical reach of ICRL.

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.

In-Context Reinforcement Learning From Suboptimal Historical Data

TL;DR

This work tackles in-context reinforcement learning (ICRL) from suboptimal historical data by introducing the Decision Importance Transformer (DIT). DIT pretrains a task-conditioned transformer on trajectories from diverse RL tasks and uses in-context advantage estimation to weight actions during a weighted maximum likelihood objective, enabling policy improvement without requiring optimal action labels. A key contribution is the in-context estimator setup—two transformers modelled to output in-context and , producing used to form exponential weights for pretraining. Empirically, DIT matches or outperforms related approaches on both bandit and multi-task MDp benchmarks, demonstrating strong generalization to unseen tasks and robustness to suboptimal data, thus broadening the practical reach of ICRL.

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.
Paper Structure (41 sections, 2 theorems, 31 equations, 8 figures, 2 algorithms)

This paper contains 41 sections, 2 theorems, 31 equations, 8 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminary
  4. Decision Importance Transformer
  5. Weighted Maximum Likelihood Estimation
  6. In-context Task Identification and Advantage Function Estimation
  7. Experiments
  8. Bandit Problems
  9. MDP Problems
  10. Discussion
  11. Results on Meta-World and Half-Cheetah
  12. Baselines
  13. Bandit Algorithms
  14. Empirical Mean (EMP).
  15. Upper Confidence Bound (UCB).
  16. ...and 26 more sections

Key Result

Proposition 4.1

Consider the following optimization problem where $\mathbb{E}_{\tau,s,a}$ is defined as in Equation eqn:loss-dit except that $a \sim \pi(a|s;\tau)$, i.e., the action is sampled from the task-conditioned policy rather than the behavioral policies: where $D_{\mathrm{KL}}$ is the Kullback–Leibler (KL) divergence, and let $\pi^\star \in \mathop{\mathrm{arg\,max}}\limits_{\pi}J(\pi)$ be its optimizer.

Figures (8)

  • Figure 1: (a) and (c) Comparison between Offline RL and ICRL. Standard offline RL trains and tests a policy $\pi$ in the same task (Env 0); ICRL pretrains TMs on trajectories collected from a family of different RL tasks (Env 1, Env 2, …, Env M), and deploys the pretrained TMs to unseen tasks (Env M+1). ICRL Deployment. The pretrained TMs generate actions conditioned on the current states and context datasets consisting of offline trajectories collected by (suboptimal) behavioral policies from the unseen tasks. (b) Supervised Pretraining. Presented with offline trajectories and optimal action labels, TMs are pretrained to predict the optimal actions for query states across RL tasks. (c) ICRL from Suboptimal Historical Data. This work addresses the challenging problem of ICRL without optimal action labels. (d) Schematic Overview of the Proposed Framework DIT. Lack of the optimal action labels, the proposed framework employs in-trajectory state-action pairs as query states and pseudo-optimal action labels, and a weighted pretraining objective, where the weights are based on the optimality of actions, estimated by a TM-based in-context advantage function estimator.
  • Figure 2: Results for Linear Bandits (lower values indicate better performance). Left: Online testing. Middle: Offline testing conditioned on trajectories gathered by highly suboptimal, randomly generated policies. Right: Offline testing conditioned on trajectories gathered by experts.
  • Figure 3: Results on Dark Room (higher values indicate better performance). (a): Change in return of policies with additional online episodes for (in-context) learning. (b) and (c): Offline evaluations with context trajectories sampled from random and expert policies.
  • Figure 4: Ablation Study on Miniworld. From left to right: online testing, offline random, offline expert.
  • Figure 5: Top row shows results on Meta-World; bottom row shows results on Half-Cheetah.
  • ...and 3 more figures

Theorems & Definitions (4)

  • Proposition 4.1
  • Proposition 4.2: Policy Improvement
  • proof : Proof of Proposition \ref{['lemma:equivalence']}
  • proof : Proof of Proposition \ref{['prop:policy-improvement']}