In-context Reinforcement Learning with Algorithm Distillation

Michael Laskin; Luyu Wang; Junhyuk Oh; Emilio Parisotto; Stephen Spencer; Richie Steigerwald; DJ Strouse; Steven Hansen; Angelos Filos; Ethan Brooks; Maxime Gazeau; Himanshu Sahni; Satinder Singh; Volodymyr Mnih

Source-linked AI summary

In-context Reinforcement Learning with Algorithm Distillation

Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, Maxime Gazeau, Himanshu Sahni, Satinder Singh, Volodymyr Mnih

arXiv:2210.14215v1Oct 25, 2022cs.LGcs.AI

TL;DR

Existing policy-distillation methods extract policies from offline RL data but do not improve through in-context trial and error. Algorithm Distillation models RL training histories with a causal transformer, enabling entirely in-context policy improvement. Across partially observed environments, AD demonstrates exploration, credit assignment, generalization, and more data-efficient learning than its source algorithm.

Problem
Policy Distillation methods learn policies from offline data but do not learn reinforcement learning algorithms that improve incrementally through in-context trial and error.
Method
Algorithm Distillation trains a causal transformer to predict actions from across-episodic learning histories generated by an RL algorithm, thereby modeling policy improvement.
Results
AD can reinforcement learn new tasks entirely in context and demonstrates exploration, temporal credit assignment, generalization, and greater data efficiency than the source algorithm.
Takeaways & Limitations
Learning histories can be distilled into an in-context RL algorithm that improves without updating transformer parameters.
Takeaways & Limitations
AD is constrained by long episodes, which require more powerful long-horizon sequential models to represent multi-episodic context.

Abstract

from arXiv · show

We propose Algorithm Distillation (AD), a method for distilling reinforcement learning (RL) algorithms into neural networks by modeling their training histories with a causal sequence model. Algorithm Distillation treats learning to reinforcement learn as an across-episode sequential prediction problem. A dataset of learning histories is generated by a source RL algorithm, and then a causal transformer is trained by autoregressively predicting actions given their preceding learning histories as context. Unlike sequential policy prediction architectures that distill post-learning or expert sequences, AD is able to improve its policy entirely in-context without updating its network parameters. We demonstrate that AD can reinforcement learn in-context in a variety of environments with sparse rewards, combinatorial task structure, and pixel-based observations, and find that AD learns a more data-efficient RL algorithm than the one that generated the source data.

1 INTRODUCTION

Offline Policy Distillation extracts policies from offline RL data but does not improve through in-context trial and error. Algorithm Distillation instead models learning histories across episodes so a transformer can learn policy improvement and reinforce learn entirely in context.

Prior work: Offline Policy Distillation uses sequence models to extract policies from offline RL data through action prediction, return conditioning, or filtering suboptimal data.This family includes single-task, same-domain multi-task, and cross-domain policy distillation approaches.
Limitations of prior work: PD policies do not incrementally improve from additional environment interaction because they learn policies rather than reinforcement learning algorithms.MGDT adapts through weight finetuning, while Gato uses expert demonstrations for task adaptation.
Motivation: Existing training data either contains no learning or uses contexts too short to capture policy improvement.Fixed expert data lacks learning progress, whereas replay-buffer data may not expose improvement across sufficiently long histories.
Core idea: A sufficiently long causal-transformer context can represent a policy improvement operator by attending to states, actions, and rewards from previous episodes.This motivates distilling an RL algorithm into an in-context RL algorithm through imitation learning.
Algorithm Distillation: Algorithm Distillation trains on many RL training histories and causally predicts actions from preceding across-episodic histories.Because the source policy improves during training, accurate action prediction requires modeling that improvement rather than only a fixed policy.
Results: AD can reinforce learn new tasks entirely in context and demonstrates exploration, temporal credit assignment, generalization, and improved data efficiency over its source algorithm.The evaluation includes partially observed environments, including pixel-based Watermaze from DMLab.

2 BACKGROUND

The paper studies reinforcement learning when agents receive partial observations rather than full states, using sequential modeling to predict actions from offline interaction data. Transformers provide the sequence-modeling mechanism, while Policy Distillation extracts policies without learning value functions.

Partially Observable Markov Decision Processes: In a POMDP, the agent receives observations with only partial information about the environment’s true state.The underlying MDP includes states, actions, rewards, discounting, and transition dynamics, but the agent acts from observations in the POMDP setting.
Partial observability: Missing goals, pixel-based observations, or both can make full state information incomplete and require inferring task information through rewards and memory.These conditions motivate environments where history is important for solving tasks.
Online and Offline Reinforcement Learning: Reinforcement learning algorithms broadly maximize cumulative discounted rewards through interaction with the environment.The background distinguishes on-policy methods, which use Monte-Carlo returns, from off-policy methods, which learn value functions approximating future return.
Self-Attention and Transformers: Self-attention projects input data into queries, keys, and values, then uses their interactions to process sequential data.Transformers commonly use self-supervised objectives that predict tokens in sequences.
Offline Policy Distillation: Offline Policy Distillation treats offline RL as sequential prediction, extracting policies by behavior-cloning action data rather than learning a value function.Policy extraction may use return conditioning or filtering of suboptimal data.
In-Context Learning: In-context learning infers tasks from context, while in-weights learning changes model parameters.The paper uses these terms to distinguish adaptation through context from adaptation through parameter updates.

3 METHOD

Algorithm Distillation models reinforcement-learning algorithms as long history-conditioned policies. It trains a causal sequence model on multi-task learning histories so fixed parameters can reproduce policy improvement through in-context action prediction.

An agent’s actions are modeled as a function of its past observations, actions, and rewards, called its history.
3.1 ALGORITHM DISTILLATION: Algorithm Distillation treats a long history-conditioned policy as an algorithm whose task-specific learning histories are generated by the source algorithm.The environment or task is denoted M, and the source algorithm P generates each task’s observation, action, and reward sequence.
3.1 ALGORITHM DISTILLATION: A causal sequence model is trained with negative log likelihood to predict actions from preceding multi-episodic histories.Training samples randomly selected multi-episodic subsequences, autoregressively predicts actions, and updates transformer parameters by backpropagation.
3.1 ALGORITHM DISTILLATION: Because source policies improve during their learning histories, accurate action prediction requires AD to infer both the current policy and its improvement operator.The context must be long enough to span learning updates, such as updates occurring across episodes.
3.2 PRACTICAL IMPLEMENTATION: AD collects learning histories by training individual gradient-based RL algorithms on randomly sampled tasks, making the data-generation step algorithm agnostic.The paper demonstrates distillation of UCB, an on-policy actor-critic, and an off-policy DQN in distributed and single-stream settings.
3.2 PRACTICAL IMPLEMENTATION: During evaluation, the pretrained transformer unrolls in an environment while storing transitions in a context queue and predicting actions without parameter updates.Across-episodic subsequences are used because causal-transformer training and inference are quadratic in sequence length.

4 EXPERIMENTAL SETUP

The experiments test in-context reinforcement learning on environments requiring memory, exploration, and task inference, using policy-distillation and online meta-RL baselines. Evaluation includes adversarial bandits, combinatorial partially observed tasks, and pixel-based Watermaze, with aggregate multi-seed returns measured against source algorithms.

4.1 ENVIRONMENTS: The evaluation environments are selected to prevent zero-shot task inference and require many tasks with episodes short enough for across-episodic transformers.The study focuses on environments that cannot be solved through zero-shot generalization after pre-training.
4.1 ENVIRONMENTS: The Adversarial Bandit has 10 arms and 100 trials, with reward concentrated on odd arms during training and even arms during evaluation.Reward is distributed under odd arms 95% of the time during training and under even arms 95% of the time during evaluation.
4.1 ENVIRONMENTS: DMLab Watermaze is a partially observable 3D visual task with 72 × 96 × 3 pixel observations, eight actions, and a randomly spawned trap-door.Maze-wall color patterns can be used to remember the goal location, and episodes last 50 steps.
4.2 BASELINES: AD is compared with Expert Distillation, the gradient-based source algorithm, and RL2, while RL2 is treated as an approximate upper bound because it trains through environment interaction.Expert Distillation uses expert trajectories only, whereas AD uses learning histories.
4.3 EVALUATION: Evaluation measures return after pretraining, with learning occurring entirely in context as the transformer interacts with each fixed test task without updating its parameters.Results average five training seeds and 20 evaluation seeds, totaling 100 seeds, and plot performance against test-time environment steps.
4.3 EVALUATION: Across the main environments, AD consistently reinforcement learns in context and is more data-efficient than the A3C or DQN source algorithm it distills.Figure 4 reports mean return ± 1 standard deviation over five training seeds with 20 test seeds each.

5 EXPERIMENTS

Experiments test whether AD can reinforcement learn in-context across bandits and challenging sparse-reward environments, and examine demonstrations, data efficiency, context length, and model capacity. AD generally learns in-context, can outperform source-algorithm data efficiency, and improves prompted policies toward optimality.

AD reliably learns in-context across Dark Room, Dark Key-to-Door, and Watermaze, while ED fails in most settings; AD matches asymptotic RL2 on Dark environments and approaches it within 13% on Watermaze.
AD explores sparse-reward tasks, performs credit assignment from single-timestep rewards, and generalizes near-optimally to mostly unseen combinatorial tasks.
AD trained on subsampled single-stream A3C histories learns a more data-efficient in-context RL algorithm than the source algorithm.The histories were subsampled by taking every 10th episode.
AD learns to improve prompted policies in-context until they are optimal or nearly optimal, whereas ED largely maintains the input policy.
Multi-episodic contexts of 2-4 episodes are necessary for near-optimal in-context RL, with initial signs emerging around one episode.

6 RELATED WORK

Related work places AD among offline policy-distillation and meta-RL methods while distinguishing its across-episodic, incremental in-context learning objective. Unlike prior approaches, AD models learning progress from RL histories rather than only fixed policies or small contexts.

Policy-distillation methods use transformers to learn single-task or multi-task policies from offline interaction data.
AD is an offline in-context meta-RL method that adapts without updating network parameters, unlike online or optimization-based meta-RL approaches.
Incremental in-context learning means improving from one’s own trial-and-error behavior, a less common capability than learning from demonstrations.
Prior work demonstrated incremental in-context learning for hyperparameter optimization by treating it as sequential prediction with a score function.

7 CONCLUSION

The paper concludes that AD distills an in-weights RL algorithm into an in-context algorithm and can learn more data-efficient algorithms than its source data generator. Its main limitation is the difficulty of modeling long multi-episodic contexts in environments with long episodes.

AD distills an in-weights RL algorithm into an in-context RL algorithm by modeling RL learning histories with a causal transformer.
AD can learn more data-efficient algorithms than those that generated its source data.
The main limitation is that long-episode environments require more powerful long-horizon sequential models to represent multi-episodic context.
The evaluation targets multi-task environments requiring exploration and feasible across-episodic modeling, constraining the considered settings.

B CLOSELY RELATED PRIOR METHODS

AD differs from closely related methods mainly through its large across-episodic context, which captures task information and learning progress. Comparisons with ED and model-capacity results clarify why fixed expert data and insufficient capacity or context limit alternatives.

Expert Distillation: ED models expert trajectories with the same across-episode architecture as AD, but its expert-only data does not show policy improvement.
Decision Transformer: DT and related policy-distillation methods use contexts too small to capture across-episode learning progress and task information; return conditioning alone is insufficient here.
Trajectory Transformer: Trajectory Transformer additionally learns a world model and uses beam search, but its smaller within-episode contexts create the same task-identification problem.
Algorithm Distillation: AD’s across-episodic context is large enough to capture learning progress and task information, distinguishing it from prior methods.
Meta Reinforcement Learning: AD is trained offline with supervised imitation, whereas prior in-context meta-RL methods such as RL2 are trained online with multi-episodic value functions and TD learning.
Model Size: Increasing model depth and embedding width improves AD performance on Dark Key-to-Door, while in-context RL emerges across all investigated model sizes.

E SOURCE ALGORITHM TRAINING RUNS

The source algorithms generate the learning histories on which Algorithm Distillation is trained. These histories capture the asymptotic performance of A3C and a Q-λ variant of DQN in Dark and Watermaze environments.

A3C and a Q-λ variant of DQN are used as source RL algorithms for Dark and Watermaze environments.
The plotted curves show the learning histories on which AD is trained.
The source-algorithm curves in Figure 4 are the same as those shown in these plots.

F LABEL SMOOTHING ABLATION

Label smoothing improves AD’s in-context learning on Dark Room (Hard) up to a point. Performance continues increasing with evaluation episodes, and the ablation averages each curve over five training seeds.

Adding label smoothing improves AD’s in-context learning ability on Dark Room (Hard).
Performance increases continually with the number of evaluation episodes when label smoothing is applied up to an effective level.
Figure 10 compares three label-smoothing α values with label smoothing turned off, averaging each curve over five training seeds.

H NUMBER OF TRAINING TASKS IN SOURCE DATA

The experiments examine how the amount and compression of source-task data affect AD. AD can learn faster than its source algorithm even from subsampled single-stream histories, while masking values of 0.3–0.5 perform best in one domain.

Dark Key-to-Door models are evaluated on a fixed test-task set for 300 episodes.
Models trained on 1, 9, or 18 Dark Key-to-Door tasks did not show in-context learning on held-out test tasks.
AD trained on 10x-subsampled A3C histories learns much faster than the source algorithm on Dark Key-to-Door test tasks.The histories retain every 10th episode, producing 200-episode compressed trajectories from 2000-episode A3C runs.
Random-masking values of 0.3–0.5 perform best on the 9x9 Dark Key-to-Door domain, with 0.3 used in all experiments.

K AD NETWORK ARCHITECTURE: TRANSFORMER VS LSTM

AD supports in-context learning with both Transformer and LSTM backbones, but the Transformer variant consistently performs better on Dark Key-to-Door and is used in the remaining experiments. Additional comparisons examine evaluation length and attention over the learning context.

Both Transformer AD and LSTM AD achieve in-context learning on Dark Key-to-Door, so AD’s success is not tied to one architecture.
The Transformer variant consistently outperforms the LSTM variant on Dark Key-to-Door and is used for the other experiments.
Figure 14 compares Transformer and LSTM AD over 300 evaluation episodes, equivalent to 15k environment steps, using five training and 20 evaluation seeds.
Figure 15 compares AD attention after 200 time-steps, when context is initially filled, with attention after 1900 steps across 38 episodes.White and gray encode attention intensity; red and blue mark episode restarts and positive-reward tokens.