Efficiently Robust In-Context Reinforcement Learning with Adversarial Generalization and Adaptation

Juncheng Dong, Hao-Lun Hsu, Miroslav Pajic, Vahid Tarokh

Keywords: in-context learning, transformers, in-context reinforcement learning, sequential decision-making

Abstract: Transformer models (TMs) pretrained on diverse datasets exhibit impressive in-context learning (ICL) capabilities, enabling them to adapt to new tasks without parameter updates. In reinforcement learning (RL), in-context RL (ICRL) leverages this capability by pretraining TMs on diverse RL tasks to adapt to unseen ones. However, the robustness of pretrained TMs in ICRL remains underexplored, with performance often degrading under disturbances in deployment. To address this, we propose a pretraining framework that augments the data with adversarial variations of training environments. This approach improves robustness and enhances generalization by increasing task diversity, while avoiding the computational overhead and pessimism of worst-case optimization used in robust RL. We further introduce an adaptive variant that uses the ICRL capability of pretrained TMs to efficiently generate high-quality data through online rollouts. Extensive experiments on a diverse set of ICRL tasks prove the efficacy of the proposed methods.

Submission Number: 13