Just-In-Time Reinforcement Learning: Continual Learning in LLM Agents Without Gradient Updates

TL;DR

Just-In-Time Reinforcement Learning (JitRL) enables continual on-the-fly adaptation of frozen LLM agents without gradient updates by storing experiences in a non-parametric memory and performing test-time policy optimization via a logit-space, KL-constrained update. The approach combines memory-based value estimation, optimistic action handling, and a closed-form logit update that is provably optimal under KL constraints, guaranteeing convergence to the KL-regularized policy. Empirically, JitRL achieves state-of-the-art results among training-free baselines and outperforms expensive fine-tuning methods on WebArena and Jericho while reducing monetary costs by over 30×. The framework generalizes across backbones and tasks, with qualitative analyses showing correction of semantic priors and strong dependence on cross-task memory for transfer. Overall, JitRL offers a practical, scalable path toward continual learning for AI agents without gradient updates, suitable for real-world deployment.

Abstract

While Large Language Model (LLM) agents excel at general tasks, they inherently struggle with continual adaptation due to the frozen weights after deployment. Conventional reinforcement learning (RL) offers a solution but incurs prohibitive computational costs and the risk of catastrophic forgetting. We introduce Just-In-Time Reinforcement Learning (JitRL), a training-free framework that enables test-time policy optimization without any gradient updates. JitRL maintains a dynamic, non-parametric memory of experiences and retrieves relevant trajectories to estimate action advantages on-the-fly. These estimates are then used to directly modulate the LLM's output logits. We theoretically prove that this additive update rule is the exact closed-form solution to the KL-constrained policy optimization objective. Extensive experiments on WebArena and Jericho demonstrate that JitRL establishes a new state-of-the-art among training-free methods. Crucially, JitRL outperforms the performance of computationally expensive fine-tuning methods (e.g., WebRL) while reducing monetary costs by over 30 times, offering a scalable path for continual learning agents. The code is available at https://github.com/liushiliushi/JitRL.

Figures (5)

• Figure 1: While standard RL performs policy gradient updates during training using previous trajectories, JitRL operates at test time. Specifically, it retrieves trajectories relevant to the current state to estimate advantages $A$, subsequently refining the output logits through a KL-regularized policy optimization objective.
• Figure 2: Overview of the Just-In-Time Reinforcement Learning (JitRL) framework. The system operates in a continuous loop: (1) In the Inference (top), the agent retrieves relevant past experiences $\mathcal{N}(s)$ from the non-parametric memory $\mathcal{M}$. The base LLM's logits $z$ are then adjusted in closed-form ($z' = z + \beta \widehat{A}$) using the estimated advantage $A$ derived from historical returns, enabling test-time policy improvement without gradient updates. (2) In the Memory Update (bottom), completed trajectories are analyzed by an evaluator to compute discounted returns $G_t$. These new experiences are stored back into $\mathcal{M}$, allowing the agent to evolve its policy across episodes.
• Figure 3: Learning curves on Jericho games. JitRL shows consistent improvement across episodes.
• Figure 4: Impact of Retrieval Neighbor Count ($k$).
• Figure 5: Performance Comparison of Text-based and Embedding-based State Representations.

Theorems & Definitions (7)

• Theorem 4.1: Optimality of Policy Update
• Theorem 4.2: Consistency of $\widehat{V}_t$, $\widehat{Q}_t$, and $\widehat{A}_t$
• Theorem 4.3: Consistency of Policy Update
• Theorem 3.7: Tracking consistency of $\widehat{V}_t$, $\widehat{Q}_t$, and $\widehat{A}_t$
• proof
• Theorem 4.1: Consistency of the closed-form policy update
• proof