Context management reduces token use in reinforcement learning agents
Downstream-Aware Context Selection for Online In-Context Reinforcement Learning
Machine Learning
Summary
Large language model agents learn from their past actions to adapt without changing their software, but remembering everything can be costly. The authors designed a method to decide which parts of the past to remember, keeping only the important history for each new decision. This saves computing resources while still letting the agent perform well in tasks like driving in simulations and exploring virtual worlds. Their approach cut the amount of information needed by about a quarter to half in different tests.
What this means in practice
- •For autonomous vehicle teams: Manage memory for reinforcement learning agents controlling vehicles to reduce computational costs while maintaining driving performance in simulation.
- •For virtual assistant developers: Efficiently select historical interaction data for language model agents to improve adaptation in ongoing conversations without increasing compute costs.
Authors
Ruihan A. Li, Shangtong Zhang, Rohan Chandra
Abstract
In-context reinforcement learning (ICRL) enables large language model agents to adapt to new environments using their interaction history without updating model parameters. However, repeatedly conditioning on growing histories can lead to substantial token cost. We propose a bounded-history context-management framework that predicts the task-dependent downstream effect of removing historical interactions to guide history selection and determine a decision-dependent context budget. Formally, our framework uses the full rolling history as a reference. The predictor evaluates removal effects, defines a deletion ordering, and applies a shared selection criterion to determine how much history to retain at each decision. We evaluate the method in closed-loop SUMO driving under held-out in-distribution, unseen-domain, and unseen-route settings, and in ScienceWorld under a continual ICRL protocol. Relative to a baseline using the full context, our method reduces total token usage by 25.7%, 25.8%, and 23.2% across the three driving settings while maintaining comparable closed-loop driving performance. In ScienceWorld, it reduces total token usage by 52.1% compared to full context and uses 30.2% and 37.8% fewer tokens than the Recent and Similarity baselines, respectively, while maintaining performance.