Robust In-Context Reinforcement Learning Under Reward Poisoning Attacks

We study the corruption-robustness of in-context reinforcement learning (ICRL), focusing on the Decision-Pretrained Transformer (DPT, Lee et al., 2023). To address the challenge of reward poisoning attacks targeting the DPT, we propose a novel adversarial training framework, called Adversarially Trained DPT (AT-DPT). Our method simultaneously trains a population of attackers to minimize the true reward of the DPT by poisoning environment rewards, and a DPT model to infer optimal actions from the poisoned data. We evaluate the effectiveness of our approach against standard bandit algorithms, including robust baselines designed to handle reward contamination. Our results show that AT-DPT significantly outperforms them in bandit settings under a learned attacker, and generalizes to more complex environments such as adaptive attackers and MDPs. It shows promise in ICRL as a meta-RL approach to learning effective corruption-robust algorithms.

Machine Learning (ML) models acting in the world can be poisoned (for example, by altering the rewards the algorithm observes) to perform malicious actions or provoke harmful responses. We design a method to make one type of ML model (the Decision-Pretrained Transformer (DPT), an in-context reinforcement learning method) robust against this type of poisoning. Our method, called Adversarially-Trained DPT (AT-DPT), simulates poisoning strategies and trains on them. This work shows promise for ensuring that AI systems which learn from rewards can still be trusted when that reward is being deliberately corrupted.