🤖 machine learning

Robust In-Context Reinforcement Learning Under Reward Poisoning Attacks

This paper proposes Adversarially Trained DPT (AT-DPT), a novel framework that simultaneously trains a Decision-Pretrained Transformer and a population of attackers to achieve robust in-context reinforcement learning against reward poisoning attacks, demonstrating superior performance over standard baselines in both bandit and complex MDP environments.

Original authors: Paulius Sasnauskas, Yiğit Yalın, Goran Radanović

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Paulius Sasnauskas, Yiğit Yalın, Goran Radanović

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Robot to Learn from a Cheat Sheet

Imagine you are teaching a robot to play a video game. Usually, you might program it with specific rules or let it practice for hours until it figures out the game.

But this paper is about a different kind of robot: one that uses In-Context Reinforcement Learning (ICRL). Think of this robot as a super-smart student who doesn't need to be reprogrammed. Instead, you just hand it a "cheat sheet" (a history of past moves and rewards) and say, "Here's how the game works; now go play." The robot reads the cheat sheet, figures out the pattern, and starts playing immediately.

The paper focuses on a specific type of robot called DPT (Decision-Pretrained Transformer). It's great at reading these cheat sheets. But, there's a problem: What if someone tampers with the cheat sheet?

The Problem: The "Poisoned" Cheat Sheet

In the real world, a bad actor (an attacker) might try to trick the robot. They can't change the robot's brain, but they can change the rewards the robot sees in its history.

The Scenario: Imagine the robot is learning to pick the best fruit in a market.
The Attack: A saboteur secretly changes the price tags on the fruit in the robot's history. They make the rotten apples look expensive (high reward) and the fresh apples look cheap (low reward).
The Result: The robot reads the poisoned history, thinks the rotten apples are the best choice, and starts buying them. This is called a Reward Poisoning Attack.

Most previous research focused on protecting robots that learn during training. This paper asks: How do we protect a robot that is learning "on the fly" just by reading a history of events?

The Solution: The "Sparring Partner" Method (AT-DPT)

The authors created a new method called AT-DPT (Adversarially Trained DPT). They didn't just try to patch the robot; they taught it how to fight back using a technique called Adversarial Training.

Think of this like a martial arts dojo:

The Student (The Robot): We want the robot to learn to pick the best actions even when the data is messy.
The Sparring Partner (The Attacker): The researchers created a "villain" AI whose only job is to try to trick the robot. This villain constantly tries to rewrite the history (the rewards) to make the robot make bad choices.
The Training Loop:
- The Villain tries to poison the robot's history.
- The Robot tries to ignore the poison and figure out the truth.
- They do this over and over again.
- Eventually, the Robot becomes so good at spotting the lies that it can still win, even when the Villain is trying its hardest to trick it.

The result is a robot that has "seen it all." It has learned that sometimes the reward numbers are lies, and it knows how to look past them to find the real truth.

How They Tested It

The researchers tested this "super-robot" in three different scenarios, like levels in a video game:

The Slot Machine (Bandits): Imagine a row of 5 slot machines. The robot has to figure out which one pays out the most. The attacker tries to lie about how much each machine pays.
- Result: The AT-DPT robot figured out the real winners, while standard robots (like Thompson Sampling or UCB) got confused and kept picking the losers.
The Maze (MDPs): Imagine a robot navigating a dark room to find a treasure. The attacker tries to lie about how good a move was.
- Result: The AT-DPT robot found the treasure much more reliably than other methods, even when the attacker was smart and adaptive (changing its lies based on what the robot did).
The Visual Maze: They even tested it on a 3D environment where the robot had to navigate using images. The AT-DPT robot still performed better than the others.

Why This Matters

The paper claims that by training the robot to expect lies (poisoned data), it becomes much more robust.

Standard robots are like people who believe everything they read in a newspaper. If the newspaper prints a lie, they believe it.
The AT-DPT robot is like a detective who has been trained by a master forger. It knows that the newspaper might be lying, so it double-checks the facts and still finds the truth.

The authors conclude that this approach (using a population of attackers to train the learner) is a powerful way to build AI that is safe and reliable, even when the data it relies on is being tampered with by malicious actors. They also noted that this method works well even if the robot isn't perfect at spotting the lies, as long as it has been trained against a variety of tricky attackers.

Technical Summary: Robust In-Context Reinforcement Learning Under Reward Poisoning Attacks

Problem Definition

This paper addresses the vulnerability of In-Context Reinforcement Learning (ICRL) to reward poisoning attacks. While traditional RL literature has extensively studied reward poisoning as a training-time attack (where an agent's policy is corrupted during learning), ICRL agents operate differently. They implement a learning algorithm "in-context" by processing a history of interactions (states, actions, rewards) within a transformer model (specifically the Decision-Pretrained Transformer, or DPT).

In this setting, an adversary can poison the test-time context by corrupting the rewards observed by the agent. Since the agent infers its policy based on this context, corrupting the rewards at test time allows an attacker to influence the agent's behavior without altering the agent's pre-trained parameters. The paper formalizes this as a test-time reward poisoning attack where the attacker modifies a fraction ( $\epsilon$ ) of the observed rewards to minimize the agent's true return, subject to a soft budget constraint.

Methodology: Adversarially Trained DPT (AT-DPT)

The authors propose AT-DPT, a novel framework that combines in-context learning with adversarial training to produce agents robust against such attacks. The core idea is to train the agent and a population of attackers simultaneously in a minimax game.

1. Attack Model

The attacker is modeled as a function $\pi^\dagger_\phi$ that observes the state, action, true reward, and the agent's in-context dataset. It generates a corrupted reward $r^\dagger$ with probability $\epsilon$ . The attacker's objective is to minimize the agent's expected return while adhering to a soft budget constraint on the mean and variance of the corrupted rewards relative to the true environment rewards. This is formulated as a penalty term in the attacker's loss function.

2. Training Procedure

The training process (illustrated in Figure 1 and Algorithm 1) involves three phases:

Pretraining: The DPT is initialized and trained via supervised learning to predict optimal actions from clean context data, similar to the original DPT approach.
Adversarial Training: This is the novel contribution. The agent ( $\pi_\theta$ $π_{θ}$ ) and a population of attackers ( $\pi^\dagger_{\phi_i}$ $π_{ϕ_{i}}^{†}$ ) are optimized concurrently:
- Attacker Update: For each task in a batch, the attacker generates corrupted trajectories. The attacker parameters $\phi$ are updated using the REINFORCE algorithm to maximize the agent's regret (minimize return) under the budget constraints.
- Agent Update: The agent is trained to predict the optimal action (provided by an oracle) given the corrupted context $D^\dagger$ collected during the interaction.
- This process repeats for $N$ rounds, allowing the agent to learn a "correction" mechanism that infers optimal actions despite poisoned inputs.
In-Context Learning (Testing): The trained agent is deployed in new environments with poisoned rewards. It updates its context dynamically at every timestep to support robustness against adaptive attackers.

The framework supports both non-adaptive attackers (which do not use the agent's history) and adaptive attackers (which use a transformer architecture to leverage the agent's past interactions).

Key Contributions

Novel Attack Modality: The paper is the first to study reward poisoning in the context of meta-RL (learning a learning algorithm). It formalizes a test-time attack where the adversary corrupts the context to influence the in-context learning process, distinct from traditional training-time policy corruption.
AT-DPT Framework: The authors introduce a practical adversarial training protocol that simultaneously optimizes a population of attackers and the victim agent. This enables the agent to learn a corruption-robust learning algorithm in-context.
Empirical Validation: The study provides a systematic evaluation across multi-armed bandits, linear bandits, and MDPs (including the Darkroom2 and Miniworld environments), comparing AT-DPT against standard algorithms (TS, UCB, PPO, Q-learning) and robust baselines (crUCB, RTS).

Experimental Results

The experiments demonstrate that AT-DPT significantly outperforms existing baselines under various poisoning levels ( $\epsilon = 0.1, 0.2, 0.4$ ) and attack budgets.

Bandit Settings: AT-DPT achieves the lowest cumulative regret compared to robust baselines like crUCB and RTS, even when the attacker is specifically trained to target AT-DPT. It also shows strong generalization, performing well against attackers trained for other algorithms (e.g., TS).
Adaptive vs. Non-Adaptive: The results indicate a trade-off; AT-DPT trained against adaptive attackers performs better against adaptive threats but slightly worse against non-adaptive ones compared to a model trained only on non-adaptive threats. However, in all cases, it outperforms standard baselines.
MDP Settings: In the Darkroom2 and Miniworld environments, AT-DPT maintains high episode rewards under attack, whereas standard RL algorithms (NPG, Q-learning) and frozen DPT models suffer significant performance degradation. Notably, AT-DPT generalizes to larger, more complex environments, suggesting scalability.
Suboptimal Demonstrations: The paper notes that AT-DPT can tolerate a small percentage (e.g., 10-30%) of suboptimal demonstrations in the training data, though performance degrades as this percentage increases.

Significance and Claims

The paper claims that its primary significance lies in establishing a principled approach to designing learning algorithms that are inherently robust to data corruptions.

Meta-RL Perspective: Unlike prior works that focus on learning a specific policy for a single environment, this work focuses on learning the algorithm itself. By utilizing a prior over the distribution of tasks, the agent can distinguish between plausible and implausible corruption schemes, effectively filtering out adversarial noise.
Generalization: The results suggest that transformer-based policies can implement algorithms that are robust against data contamination, offering a promising direction for safe deployment of ICRL agents (such as LLM-based agents) in environments where context may be compromised (e.g., via RAG mechanisms).
Theoretical Insight: The authors posit that AT-DPT approximates a Nash equilibrium where the agent implements a "corrected" posterior sampling. As long as the underlying reward means remain identifiable under the corruption process, the agent can achieve robust performance.

The authors remain modest regarding theoretical guarantees, acknowledging that a complete convergence analysis of the bi-level optimization to an approximate equilibrium is left for future work. They also note the limitation that their current method relies on oracle-provided optimal actions for training, though they demonstrate some resilience to suboptimal demonstrations.