Source-linked AI summary
Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation
Yuanqing Ma, Zhenrui Zheng, Chenjun Xiao
TL;DR
Algorithm Distillation can perform in-context reinforcement learning, but long histories require expansive, costly context windows and truncation may remove information needed for long-horizon adaptation. RAD recursively compresses interaction histories into latent memory alongside recent transitions, and across diverse benchmarks it generally matches or exceeds AD with substantially shorter contexts.
Problem
Method
RAD uses a Compression Transformer to recursively condense interaction histories into fixed-length latent memory, which an AD Transformer combines with recent transitions when generating actions.
Source: §1 · p. 1, §4.1 · p. 4Results
RAD generally matches or exceeds AD across diverse RL benchmarks while using substantially shorter contexts.
Source: §1 · p. 1Takeaways & Limitations
The results support recurrent context compression as an effective approach to scalable long-horizon in-context reinforcement learning.
Source: §1 · p. 1Training stability is limited because recurrent memory updates require backpropagation through time, increasing optimization and computation difficulty over long horizons.
Source: §8 · p. 9
Abstract
from arXiv · showhide
Algorithm Distillation (AD) has demonstrated the remarkable ability of Transformers to perform in-context reinforcement learning without explicit weight updates. However, capturing long-term learning progress necessitates expansive context windows, which incur prohibitive memory costs and limit scalability in complex, long-horizon tasks. To address this bottleneck, we propose Recurrent Algorithm Distillation (RAD). RAD employs a dual-component architecture: a Compression Transformer that distills extended interaction histories into compact latent tokens, and an AD Transformer that auto-regressively generates actions using a hybrid context of these compressed memories and recent transitions. By maintaining a fixed-size latent buffer, RAD decouples the effective history length from computational complexity, functionally providing the model with a long-horizon memory. Empirical evaluations across diverse environments demonstrate that RAD matches the asymptotic performance of standard AD with significantly reduced context window sizes, offering a scalable solution for efficient in-context decision-making.
1 INTRODUCTION
Algorithm Distillation uses interaction histories for in-context learning, but long histories strain Transformer context windows and truncation may discard information needed for long-horizon adaptation. RAD recursively compresses history into fixed-size latent memory and generally matches or exceeds AD with substantially shorter contexts.
1 INTRODUCTION
- AD requires expansive context windows to capture learning progress, but Transformer quadratic complexity makes them costly and truncation may discard information needed for long-horizon adaptation.Source: §1 · p. 1
- RAD combines a Compression Transformer that turns extended interaction histories into fixed-length latent tokens with an AD Transformer that acts using latent memory and recent transitions.As working memory fills, the system recursively recompresses context to update long-term memory while retaining recent transitions.Source: §1 · p. 1, Figure 1 · p. 2
- RAD generally matches or exceeds AD across diverse RL benchmarks while using substantially shorter contexts.The paper presents recurrent context compression as an effective approach to scalable long-horizon in-context reinforcement learning.Source: §1 · p. 1
2 BACKGROUND
The background frames reinforcement learning as interaction within a partially observable process and describes Transformers as sequence models that use self-attention to aggregate information across tokens. Direct token-to-token dependencies let self-attention capture both nearby and distant relationships.
2 BACKGROUND
- A POMDP is defined by actions, observations, transitions, rewards, and a discount factor, with the agent choosing actions from observations and interaction history.The objective is to maximize expected return.Source: §2 · p. 2
- Transformers process sequences through self-attention, which lets each token dynamically aggregate information from other tokens.Source: §2 · p. 2
- Self-attention models pairwise token dependencies, allowing Transformers to capture both local and long-range relationships without recurrent computation.Source: §2 · p. 2
3 IN-CONTEXT RL AND ALGORITHM DISTILLATION
In-context reinforcement learning adapts from interaction history without parameter updates, and Algorithm Distillation trains a Transformer to model the action choices in RL learning histories. The paper characterizes AD as posterior sampling, while noting that fixed sliding windows can remove distal context needed for long-range inference.
3 IN-CONTEXT RL AND ALGORITHM DISTILLATION
- In-context reinforcement learning uses interaction history to adapt during inference without gradient-based parameter updates.The context window supports policy refinement in a single forward pass.Source: §3 · p. 3
- Algorithm Distillation models RL training histories as history-conditioned policies to induce the source algorithm’s policy-improvement rule and solve novel tasks without further weight updates.The histories are collected by running a source algorithm across sampled tasks, and a Transformer is optimized with a cross-entropy loss.Source: §3 · p. 3
- AD implicitly performs posterior sampling by inferring the task from history and observation, then acting according to the source algorithm for that task.Source: §3 · pp. 3–4
- Sliding-window truncation discards older interactions, leaving a potentially incomplete history that can sever long-term dependencies needed for accurate posterior inference.The paper identifies this as a scalability bottleneck for complex tasks requiring extensive historical context.Source: §3 · p. 4
4 RECURRENT ALGORITHM DISTILLATION
RAD addresses fixed-window loss of long-term history with recent-transition working memory and compressed latent memory. It recurrently updates that memory and trains the policy to predict actions from source-RL histories.
4 RECURRENT ALGORITHM DISTILLATION
- RAD pairs recent-transition working memory with compressed long-term memory to preserve historical dependencies beyond the sliding window.A transformer compressor distills context into latent tokens, and a reconstruction model is trained to recover the original sequence from them.Source: §4 · p. 4, §4.1 · p. 4
4.1 MEMORY REPRESENTATION
- Working memory retains recent transitions, while long-term memory stores latent representations of distal dependencies and updates when the working-memory capacity is reached.The long-term memory begins empty; compression is designed to preserve dependencies that would otherwise be lost as the working window advances.Source: §4.1 · p. 4
4.2 TOKENIZATION
- RAD tokenizes each observation, action, and reward separately, yielding three tokens per timestep and preserving each timestep’s tuple during compression.The latent memory uses tokens of the same dimension, and its capacity is set to a multiple of three.Source: §4.2 · pp. 4–5
4.3 INFERENCE
- When working memory reaches capacity, a compressor forms a candidate latent state, a GRU-style gated update integrates it with prior memory, and the window retains p recent transitions.Actions are generated autoregressively from both memory components; at the first compression event, the candidate directly initializes long-term memory.Source: §4.3 · p. 5
4.4 TRAINING
- Training recursively compresses source-RL histories into latent states, supervises action prediction on the final working window, and backpropagates through only the most recent G compression events.RAD uses the same source-learner history dataset as standard AD.Source: §4.4 · p. 5
5 EXPERIMENTS
RAD is evaluated on discrete and continuous tasks, including delayed bandits, grid-world mazes, and robotic manipulation. The reported results show history retention and AD-comparable or better performance, alongside reduced inference computation in most environments.
5.1 ENVIRONMENTS
- The benchmark suite spans discrete and continuous environments with varying task complexity and horizon lengths, including bandits, two grid-world mazes, and Meta-World manipulation tasks.The mazes include 81 Darkroom tasks and 6,561 Dark Key-to-Door configurations; Meta-World ML1 uses 50 task seeds and 100-step horizons.Source: §5.1 · p. 6
5.2 DATASET GENERATION
- Training histories come from UCB for bandits and PPO for MDPs; Dark Key-to-Door PPO inputs stack eight observations, while the dataset stores only the current observation.Source: §5.2 · p. 6
5.3 RESULTS
- RAD outperforms all baselines in Darkroom and exceeds AD in Dark Key-to-Door, while DPT and IDT score higher there but require domain knowledge.The figure compares average episode rewards across the two environments; results average 3 training seeds and 20 evaluation seeds.Source: Figure 3 · p. 7
- Across delayed bandits, grid-world, and Meta-World, RAD retains useful history, outperforms AD on grid-world, and generally matches or exceeds AD on continuous tasks.The cited figures report cumulative regret for delayed bandits and success rates for Meta-World; RAD has lower regret than ADlong at a 50-step delay.Source: §5.3 · pp. 7–8, Figure 4 · p. 7
- RAD reduces inference computation by more than 80% in most environments, except Dark Key-to-Door, where sparse rewards and preserving useful information are challenging.Table 1 reports inference FLOPs per action for each sequence.Source: §5.3 · p. 8, Table 1 · p. 8
6 ABLATION STUDIES
The ablations compare long-term memory update methods and tokenizer choices. GRU gating has a higher median reward and narrower interquartile range than the alternative update methods.
6 ABLATION STUDIES
- GRU gating yields a higher median reward and narrower interquartile range than the other long-term memory update methods.The authors suggest this may relate to stabilizing gradient flow during backpropagation through recurrent memory.Source: §6 · p. 8
- Figure 6 compares AD/RAD performance with different tokenizers on Dark Key-to-Door.The supplied figure caption identifies the comparison but gives no outcome.Source: Figure 6 · p. 8
7 RELATED WORK
Related work covers methods for in-context reinforcement learning and bounded-memory sequence modeling. RAD focuses on retaining extended interaction history under a fixed memory budget.
7 RELATED WORK
- Prior ICRL work changes policy learning methods, while RAD targets the interaction history retained under a fixed memory budget.The cited directions include noise, variable action spaces, larger training tasks, mixture-of-experts, reweighted losses, and richer prediction objectives.Source: §7 · p. 9
- Long-context approaches include sparse or linear attention, bounded recurrent or compressed memories, and external retrieval stores.Examples include Transformer-XL, Compressive Transformer, Block-Recurrent Transformer, and Memorizing Transformers.Source: §7 · p. 9
- In RL, RAD preserves recent interactions at full fidelity and recursively consolidates older learning history into a fixed-size latent state for action prediction.This contrasts with retrieval of selected past sub-trajectories in RA-DT and explicitly long contexts in AMAGO.Source: §7 · p. 9
8 CONCLUSION
RAD is presented as a scalable approach to long-horizon ICRL, but recurrent memory updates make training stability a limitation. The appendix gives inference and training pseudocode.
8 CONCLUSION
- Recurrent memory updates require backpropagation through time, increasing optimization and computation difficulty over long horizons.The conclusion identifies training stability as a key limitation.Source: §8 · p. 9
A RAD ALGORITHM DETAILS
- RAD inference samples actions from the policy using the latent state, working memory, and current observation, then updates memory when the working-memory capacity is reached.After compression, it retains the most recent K−p transitions; later updates combine the prior latent state with a compressed candidate.Source: A RAD ALGORITHM DETAILS · p. 14
- The training pseudocode separates compression pretraining from action distillation.The supplied steps specify sampling training sequences and updating the compression and reconstruction models before the action-distillation phase.Source: A RAD ALGORITHM DETAILS · p. 14
B EXPERIMENT DETAILS
The experiment appendix describes the model architectures, source-algorithm histories, and training configurations. It distinguishes transition-level experience from Transformer tokens and outlines the policy and compressor designs.
B EXPERIMENT DETAILS
- The appendix specifies architectures, source RL algorithms, and training configurations, and distinguishes environment transitions from Transformer tokens.Each completed transition contributes observation, action, and reward tokens; a latent memory token is a learned vector, not an individual transition in token length.Source: B EXPERIMENT DETAILS · p. 14
B.1 MODEL ARCHITECTURES
- The AD policy uses a GPT-2-style decoder-only Transformer to predict actions from temporally interleaved observation, action, and reward embeddings.The Compression Transformer appends learned query tokens, which use self-attention and cross-attention to produce candidate updates from embedded history.Source: §B.1 · p. 15
B.2 SOURCE REINFORCEMENT-LEARNING ALGORITHMS
- Offline datasets preserve the temporal order and learning progression of source-algorithm histories; UCB distractor transitions remain in input history but do not change its statistics.Only genuine arm pulls contribute action targets during distillation; PPO is used for Grid-World and Meta-World, with 100 parallel environment streams per source learner.Source: §B.2 · p. 15
B.3 TRAINING DETAILS
The training details specify the optimization setup and a staged compression curriculum for Grid-World and Meta-World. The hyperparameter tables collect architecture, context, source-algorithm, and training settings.
B.3 TRAINING DETAILS
- Training uses AdamW with weight decay 0.01, gradient clipping at an ℓ2 norm of 1.0, and linear warmup followed by cosine learning-rate decay.Experiments used eight NVIDIA GeForce RTX 4090 GPUs, with sequence models implemented in PyTorch.Source: §B.3 · pp. 15–16
- Grid-World and Meta-World training stages increase allowed compression counts and shift sampling toward longer histories.Compression counts are grouped as short (c = 0), medium (c = 1–2), long (c = 3–8), and very long (c ≥9); unavailable harder-category probability goes to the largest allowed count.Source: §B.3 · p. 16, Table 5 · p. 17
- At curriculum transitions, learning rates warm up from 10% to their group-specific peak and decay toward 10%; Delayed-Bandit uses no staged compression curriculum.Initial warmup is 2,000 steps; later warmups are 1,000 steps for Darkroom and Meta-World and 2,000 for Dark Key-to-Door.Source: §B.3 · p. 16
B.4 HYPERPARAMETERS
- Tables 2–5 collect architecture, context, source-algorithm, and training settings; Table 2 measures context windows and capacities in tokens unless specified otherwise.Table 2 notes that one transition equals three tokens; DR denotes Darkroom and DKTD denotes Dark Key-to-Door.Source: §B.4 · p. 16, Table 2 · p. 16, Table 3 · p. 17
C EXTENDED EXPERIMENTS
The extended experiments compare RAD's convergence, compressor choices, and long-term memory sizes, while the related-work section distinguishes RAD's learning-history memory from prior trajectory representations. RAD ultimately matches AD-long performance with a smaller inference context, though it converges more slowly.
C EXTENDED EXPERIMENTS
- In Darkroom, the deterministic continuous AE has the highest late-stage returns; neither VAE nor VQ-VAE improves downstream performance in this setting.VQ-VAE improves faster initially but fluctuates more and ultimately reaches returns comparable to VAE; VAE posterior means have a small late-stage advantage over sampling.Source: Figure 7 · p. 18, C EXTENDED EXPERIMENTS · p. 18
- A 15-latent-token long-term memory performs best among the tested sizes while remaining relatively small.RAD still exhibits in-context adaptation with a small long-term memory in Darkroom.Source: C EXTENDED EXPERIMENTS · p. 18, Figure 8 · p. 19
- RAD converges around 60k updates and matches AD-long performance with a significantly smaller context window, while AD-long converges to 6.9 and AD-short plateaus.AD-short descends fastest early; AD-long crosses below it around 25–30k updates, while RAD converges slowest. Figure 9 reports mean cumulative expected regret across 100 held-out tasks under delay 50.Source: C EXTENDED EXPERIMENTS · p. 18, Figure 9 · p. 19
D EXTENDED RELATED WORK
- Unlike prior trajectory representations aimed at control, prediction, or reconstruction, RAD's long-term memory summarizes experience across episodes to continue an agent's adaptation process.The related-work discussion contrasts this target with SeCTAR, history-based representations, HELM, and trajectory simplification.Source: D EXTENDED RELATED WORK · p. 20