Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation

Yuanqing Ma; Zhenrui Zheng; Chenjun Xiao

cs.LGSep 28, 2026

Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation

Authors: Yuanqing Ma, Zhenrui Zheng, Chenjun Xiao

Organizations: The Chinese University of Hong Kong, Shenzhen

Abstract

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.

Figures & tables

Environment	AD FLOPs / Action	RAD FLOPs / Action	Compute Ratio $\downarrow$
Bandit	59.8 M (short)	43.0 M	0.72 $\times$
	263.4 M (long)		0.16 $\times$
Darkroom	144.8 M	27.8 M	0.19 $\times$
Dark Key-to-Door	203.4 M	85.9 M	0.42 $\times$
Meta-World	1.84 G	130.6 M	0.07 $\times$

Table 1: Inference FLOPs per action for each sequence.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

\mathcal{L}_{\mathrm{pre}}=\mathbb{E}_{x_{t-k:t}\sim\mathcal{D}}\left[\left\|d_{\psi}\!\left(q_{\omega}(x_{t-k:t})\right)-y\right\|_{F}^{2}\right]

Appendix

Algorithm 2 RAD Training

Hyperparameter	Delayed-Bandit	DR	DKTD	Meta-World
AD baseline
Policy model dimension	64	64	64	64
Policy Transformer layers	4	4	4	4
Policy attention heads	4	4	4	8
Policy feed-forward dimension	256	256	256	256
Attention / residual dropout	0.1	0.1	0.1	0.1

Appendix

Table 2: AD and RAD hyperparameters. RAD uses the same policy architecture as AD. Context windows and capacities are measured in tokens unless otherwise specified (1 transition = 3 tokens).

Hyperparameter	DR	DKTD	Meta-World
Parallel streams per learner	100	100	100
Rollout steps per stream	80	50	100
Optimization minibatch size	40	100	200
Epochs per rollout update	20	10	20
Total Timesteps	100,000	100,000	1,000,000
Learning rate	$3\!\times\!10^{-4}$	$3\!\times\!10^{-4}$	$3\!\times\!10^{-4}$

Appendix

Table 3: PPO source-algorithm configuration.

AD policy training
Hyperparameter	Delayed-Bandit	DR	DKTD	Meta-World
Training steps	50k / 100k a	50k	50k	50k
Configured batch size	64	512	512	256
Peak learning rate	$3\!\times\!10^{-4}$	$3\!\times\!10^{-4}$	$3\!\times\!10^{-4}$	$3\!\times\!10^{-4}$
Warmup steps	500	1,000	1,000	5,000
RAD compression pretraining

Appendix

Table 4: AD/RAD training hyperparameters

Environment	Start step	Max. $c$	Short	Medium	Long	Very long
DR	0	1	0.60	0.35	0.05	0.00
DR	30k	3	0.35	0.40	0.20	0.05
DR	50k	6	0.25	0.30	0.30	0.15
DR	75k	$\infty$	0.25	0.25	0.25	0.25
DKTD	0	2	0.30	0.55	0.15	0.00
DKTD	25k	4	0.20	0.45	0.30	0.05

Appendix

Table 5: RAD curriculum. Probabilities correspond to short, medium, long, and very-long compression-count categories.

Explore similar work

LongAct: Harnessing Intrinsic Activation Patterns for Long-Context Reinforcement Learning
Apr 16, 2026Bowen Ping, Zijun Chen, Tingfeng Hui +4Efficient Long-Context Inference Large Language Model Reinforcement Learning
1Peking University · 2Shanghai Jiao Tong University · 3Beijing University of Posts and Telecommunications
Learning What to Remember: Test-Time Training via Context Distillation
Aug 3, 2026Zixuan Wang, Xingyu Dang, Rui-Jie Zhu +4Efficient Long-Context Inference Test-Time Training
Princeton University · UC Santa Cruz · Carnegie Mellon University +1
MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents
Aug 7, 2026Zhiyuan Liu, Tinghong Ye, Chenghao Liu +2Token-Level Supervision Long-Horizon Task Planning
School of Advanced Manufacturing and Robotics, Peking University · Zhejiang University · Harbin Institute of Technology

Apr 16, 2026, cs.LGJ/K move · Enter open · S save

LongAct: Harnessing Intrinsic Activation Patterns for Long-Context Reinforcement Learning

Bowen Ping, Zijun Chen, Tingfeng Hui, Qize Yu+3

1Peking University · 2Shanghai Jiao Tong University · 3Beijing University of Posts and Telecommunications

Efficient Long-Context Inference Large Language Model Reinforcement Learning Long-Context Reasoning

Open paper AskPDF arXiv

Reinforcement Learning (RL) has emerged as a critical driver for enhancing the reasoning capabilities of Large Language Models (LLMs). While recent advancements have focused on reward engineering or data synthesis, few studies exploit the model's intrinsic representation characteristics to guide the training process. In this paper, we first observe the presence of high-magnitude activations within the query and key vectors when processing long contexts. Drawing inspiration from model quantization -- which establishes the criticality of such high-magnitude activations -- and the insight that long-context reasoning inherently exhibits a sparse structure, we hypothesize that these weights serve as the pivotal drivers for effective model optimization. Based on this insight, we propose LongAct, a strategy that shifts from uniform to saliency-guided sparse updates. By selectively updating only the weights associated with these significant activations, LongAct achieves an approximate 8% improvement on LongBench v2 and enhances generalization on the RULER benchmark. Furthermore, our method exhibits remarkable universality, consistently boosting performance across diverse RL algorithms such as GRPO and DAPO. Extensive ablation studies suggest that focusing on these salient features is key to unlocking long-context potential.

Abstract

Explore similar work

LongAct: Harnessing Intrinsic Activation Patterns for Long-Context Reinforcement Learning

Learning What to Remember: Test-Time Training via Context Distillation

MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents