Can In-Context Reinforcement Learning Recover From Reward Poisoning Attacks?

TL;DR

This paper investigates the robustness of in-context reinforcement learning models against reward poisoning attacks and introduces an adversarial training method to improve their resilience in both bandit and MDP environments.

Contribution

We propose AT-DPT, a novel adversarial training framework that enhances the robustness of Decision-Pretrained Transformers against reward poisoning attacks.

Findings

AT-DPT outperforms standard and robust baselines in bandit settings.

The robustness of AT-DPT extends to complex MDP environments.

The method effectively defends against adaptive attackers.

Abstract

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 Decision-Pretrained Transformer (AT-DPT). Our method simultaneously trains an attacker 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 the proposed method significantly outperforms these baselines in bandit settings, under a learned attacker. We additionally evaluate AT-DPT on an adaptive attacker, and observe similar results. Furthermore, we extend our evaluation to the MDP setting, confirming that the robustness observed in bandit scenarios generalizes to more complex environments.