PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yibo Li, Zijie Lin, Ailin Deng, Xuan Zhang, Yufei He, Shuo Ji, Tri Cao, Bryan Hooi

arXiv:2601.18510v3cs.LGcs.AI

TL;DR

LLM agents cannot continually adapt after deployment because their weights are frozen, while conventional RL is computationally expensive and risks catastrophic forgetting. JitRL performs test-time policy optimization without gradient updates by retrieving experiences, estimating advantages, and modulating output logits. It achieves state-of-the-art performance among training-free methods, surpasses weight-update baselines, and reduces monetary costs by over 30×.

  • Problem

    Frozen deployment weights limit continual adaptation, while conventional RL requires costly updates and risks catastrophic forgetting.

  • Method

    JitRL maintains a dynamic memory of state-action-reward trajectories, retrieves relevant experiences, estimates action advantages, and directly adjusts LLM output logits without gradient updates.

  • Results

    JitRL achieves state-of-the-art performance across WebArena and Jericho, outperforming training-free and weight-update methods while reducing monetary costs by over 30×.

  • Takeaways & Limitations

    JitRL offers a scalable and efficient path for continual learning in LLM agents through inference-time policy optimization.

  • Takeaways & Limitations

    JitRL depends on actions proposed by the frozen base model and on accurate step-wise rewards from an LLM evaluator.

Abstract

from arXiv · show

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.

1. Introduction

LLM agents struggle to continually adapt because deployment freezes their weights, while conventional RL is costly and prone to catastrophic forgetting. JitRL addresses this with test-time, gradient-free policy optimization using retrieved experience and advantage-based logit updates.

  • Motivation: Frozen weights prevent current AI agents from learning continually in unfamiliar or dynamic environments.This limits practical adaptation and can cause agents to repeat errors.
  • Motivation: Conventional RL requires substantial training data and computation, frequent updates, and risks catastrophic forgetting.Prior work also finds limited improvements in some evaluation settings.
  • Motivation: In-context learning can struggle with long interaction histories and lacks RL’s ability to optimize rewards for skills difficult to express textually.These limitations motivate a more flexible test-time learning approach.
  • JitRL: JitRL stores state-action-reward trajectories, retrieves experiences relevant to the current state, estimates action advantages, and adjusts output logits without gradient updates.The method uses a dynamic memory bank to learn just in time.
  • JitRL: JitRL’s logit update is theoretically the exact closed-form solution to policy optimization under a KL constraint.This connects the direct output modification to a principled RL objective.
  • Results: JitRL achieves state-of-the-art results across WebArena and Jericho, outperforming training-free and weight-update methods while reducing conventional RL monetary costs by over 30×.The experiments also report consistent gains across LLM backbones and generalization to unseen tasks.

2. Related Work

Prior work enhances LLM agents with gradient-based RL or external memory at test time. JitRL instead treats retrieved memory as a non-parametric policy distribution and updates logits directly.

  • Reinforcement Learning: Gradient-based RL methods align LLM agents for complex tasks but are computationally expensive and produce static models that adapt poorly to distribution shifts.The related work discussion contrasts their training cost and limited adaptability with JitRL.
  • Training-Free Inference Enhancement: Recent training-free systems use external memory to store historical interactions, textual skills, failures, or dynamically indexed knowledge.Examples include MemGPT, Generative Agents, Voyager, Reflexion, and A-mem.
  • Training-Free Inference Enhancement: JitRL differs by treating memory as a non-parametric policy distribution rather than merely retrieving text for in-context learning.It performs soft updates directly on LLM logits without parameter-update overhead.

3. Preliminaries

RL optimizes a policy to maximize expected cumulative reward, while policy gradients shift probability toward higher-return actions. The advantage function evaluates actions relative to a state baseline, but estimating it can require costly value-network training.

  • RL Objective: RL optimizes a policy πθ that maps states to action distributions while maximizing expected cumulative reward.The objective is defined over trajectories and their cumulative discounted returns.
  • Policy Gradient: Policy gradients shift probability mass toward actions with higher returns by updating model parameters in the gradient direction.This is the fundamental optimization approach described in the preliminaries.
  • Advantage Function: The advantage function evaluates a selected action relative to a baseline, quantifying how much better it is than average performance in that state.It is defined as A(s, a) = Q(s, a) − V(s).
  • Advantage Function: Q(s, a) represents expected return for taking action a in state s, while V(s) represents the average expected return of being in state s.These quantities provide the action-value and state-value components of the advantage.
  • Limitation: Estimating advantages typically requires training additional value networks, making RL computationally expensive and limiting test-time adaptation.The resulting model is static rather than flexibly adapting after deployment.

4. Method

JitRL performs test-time policy optimization by retrieving experience-based value estimates from dynamic memory and using their advantages to adjust a frozen LLM without gradient updates. Its update is derived as the exact solution to a KL-constrained objective, while the value and policy estimates are theoretically shown to converge under stated assumptions.

  • Framework: JitRL modulates a frozen prior policy toward an optimal posterior at test time instead of updating model parameters.The framework constructs experience memory, estimates values through retrieval, and updates policy logits using estimated advantages.
  • Memory Construction: Completed trajectories are evaluated for step-wise rewards, converted into discounted returns, and stored as compact state-action-return triplets.The memory preserves structured states that retain task-relevant semantics while discarding irrelevant details.
  • Value Estimation: At inference, JitRL retrieves top-k neighboring transitions for the current structured state and estimates state and action values from their returns.Known actions use historical returns, while unseen actions receive an uncertainty-dependent optimistic value with probability λ or Q(s, a) = 0 otherwise.
  • Value Estimation: The test-time advantage is obtained by centering each retrieved action value against the local state baseline, identifying actions above the neighborhood average.This retrieval-based estimate serves as a proxy for the policy advantage without a trained value network.
  • Policy Update: JitRL maximizes expected advantage while constraining KL divergence from the frozen reference policy, with β controlling the constraint strength.The resulting optimal policy is implemented directly in logit space as z′(s, a) = z(s, a) + β · bA(s, a), followed by Softmax recovery.
  • Theoretical Analysis: Theoretical results establish exact optimality of the logit update, convergence of value and advantage estimators, and convergence of the policy update to the KL-regularized optimum.The estimator result applies under the paper’s stated assumptions, and the convergence claim is formulated for finite candidate action sets or uniform convergence on the candidate set.

5. Experiments

Experiments evaluate JitRL across WebArena and Jericho, comparing training-free and weight-update methods under repeated-episode learning. Results show strong benchmark performance, generalization, robust logit-based improvements, and substantially lower monetary cost.

  • Experimental Setup: Experiments compare JitRL with training-free and weight-update methods across WebArena and Jericho using repeated sequential episodes.WebArena covers realistic web navigation, while Jericho covers interactive-fiction games; WebArena tasks run for L = 5 episodes.
  • Evaluation Protocol: Average Success Rate measures overall learning efficiency, while Final Success Rate measures converged capability at the final episode.A larger Final–Avg gap indicates a steeper learning curve and greater improvement over time.
  • WebArena Results: JitRL outperforms all training-free WebArena baselines in both cumulative (Avg) and converged (Final) success rates.The largest reported gain is +73.2% over Static in Shopping, where trajectories are highly reusable.
  • WebArena Results: JitRL achieves state-of-the-art performance on held-out WebArena-Lite through inference-time optimization, despite WebRL and SFT using weight updates.WebRL and SFT were trained on the remaining WebArena tasks, while JitRL was evaluated without parameter training.
  • Jericho Results: On Jericho, JitRL achieves the highest scores across all three games, learns rapidly, widens its advantage over episodes, and converges with reduced late-stage variance.JitRL becomes competitive within the first 10–15 episodes, whereas Memory and AWM plateau early and GRPO remains highly variable.
  • Generalization: JitRL retains strong performance across different LLM backbones and unseen tasks, indicating model-agnostic behavior and transfer of abstract procedural knowledge.Cross-task evaluation restricts retrieval to memories from disjoint tasks; qualitative cases show memory correcting erroneous links and noisy searches.
  • Ablation Studies: Direct Logit Update outperforms Prompt Update on Admin and Reddit websites when both use the same retrieved memory.The authors attribute this to direct modulation of the output distribution, which avoids attention and instruction-following problems as context grows.
  • Cost Analysis: JitRL achieves superior performance at lower monetary cost than computationally intensive weight-update methods.The cost comparison contrasts API-based inference costs for JitRL with NVIDIA H200 training expenses for WebRL.

6. Conclusion

JitRL is a training-free framework for continuous test-time adaptation that directly optimizes logits in frozen LLMs. It outperforms training-free and expensive weight-update baselines while reducing monetary costs by over 30×.

  • JitRL enables frozen LLMs to continuously adapt at test time by directly optimizing logits without gradient-based updates.
  • JitRL outperforms existing training-free methods and expensive weight-update baselines while reducing monetary costs by over 30×.

Limitations

JitRL’s effectiveness is bounded by the frozen base model, reward-evaluation accuracy, and how well task-relevant information can be represented textually.

  • JitRL cannot discover actions absent from the frozen base model’s candidate set.Its action re-weighting operates only over actions proposed by the base model.
  • Inaccurate step-wise reward credit assignment can produce inaccurate advantages and degrade policy quality.JitRL relies on an LLM evaluator to generate accurate step-wise rewards.
  • Text-based state representation and retrieval may miss relevant similarities when critical task patterns are difficult to express textually.The paper gives spatial reasoning and time-series forecasting as examples where memory-based value estimation may be limited.

Impact Statement

JitRL’s theoretical and empirical results support inference-time policy adaptation through retrieved experience, while the evaluation compares it with adaptive and non-adaptive baselines across agentic settings.

  • JitRL’s additive logit update is presented as the exact solution to a KL-constrained policy optimization objective.The objective maximizes estimated advantage while constraining divergence from the reference policy.
  • Under stated regularity, coverage, noise, action-frequency, and slow-drift assumptions, the estimated values and advantages converge to their true counterparts.The consistency result uses growing memory and neighborhood conditions, state regularity, bounded zero-mean noise, and diminishing policy drift.
  • The resulting policy update converges to the KL-regularized policy defined by the true advantage values.
  • The evaluation includes static, memory-based, reflective, evolutionary, workflow-based, gradient-based, supervised, and WebRL baselines.These baselines differ in whether they retain transcripts, generate reflections, evolve configurations, reuse workflows, or update weights.

F. Implementation Details

JitRL constructs compact task-relevant states, retrieves structurally or semantically similar experiences, evaluates step-wise rewards, and uses those experiences to guide candidate action selection.

  • State Representation: JitRL compresses noisy observations into structured states that preserve task-relevant semantics for retrieval.The design maps functionally equivalent states to similar representations; WebArena additionally uses regularized URLs as structural proxies.
  • Action Handling: The augmented candidate set merges the LLM’s top-k actions with unique actions found in the retrieved neighborhood.Retrieved actions outside the LLM’s candidate set receive a neutral initial logit.
  • Memory Retrieval: JitRL retrieves experiences using task-specific matching, combining page-type filtering and Jaccard similarity for WebArena with hybrid semantic retrieval for Jericho.
  • Reflective Step-wise Rewards: An LLM evaluator assigns step-wise rewards by assessing each action’s usefulness, harm, neutrality, and certainty.Scores range from −3 to +3 and are used to quantify each action’s contribution to the task goal.
  • Comparison with WebRL: JitRL outperforms WebRL in most domains and achieves a higher overall average success rate under the compared settings.The paper attributes the gap to WebRL’s greater need for environment interactions and high-quality samples in low-data scenarios.
  • Empirical Behavior: JitRL’s memory supports non-obvious corrections in Jericho, including rewarding direct task-relevant actions, puzzle-specific commands, and prerequisite safety sequences.

M. Step Analysis

The analyses examine task complexity, representation choices, hyperparameter sensitivity, retrieval scalability, and GRPO training settings across WebArena and Jericho. Results indicate stable performance across reasonable parameter ranges and improving performance as memory grows.

  • Task complexity: Map tasks require 9.39 average steps, compared with 4.57 for Reddit and 5.55 for Admin.The differences reflect more sequential navigation operations on Map websites.
  • Hyperparameter sensitivity: Performance remains stable across reasonable exploration-rate λ and UCB-bonus α ranges in both Jericho and WebArena.WebArena favors smaller λ, while α = 5 provides the best overall balance between exploration and over-optimism.
  • Memory scalability: JitRL performance continuously improves as the memory bank grows, while retrieval overhead remains negligible relative to LLM inference time.The scalability evaluation uses the Library game over 50 episodes and reports retrieval latency and average score by memory-size bins.
  • GRPO analysis: GRPO hyperparameter sweeps evaluate rollout count, batch size, PPO epochs, and learning rate over 50 Zork1 training steps.The study records validation mean, maximum, and minimum scores across configurations and their training-episode consumption.

Q.4. Analysis

The analysis shows that tuned GRPO can approach JitRL on Zork1 but requires substantially more trajectories, while JitRL’s policy updates are illustrated as changing game-specific decisions and prioritizing high-reward actions.

  • GRPO comparison: 40.7 mean and 55 maximum validation scores are achieved by the best GRPO configuration on Zork1.The configuration uses rollout=8, batch size=8, PPO epochs=1, and lr=1e-5.
  • Policy improvement: JitRL’s qualitative policy changes override exploratory defaults, learn game-specific mechanics, and prioritize high-reward actions.The comparison contrasts base logits with logits after the memory-based update.
  • GRPO comparison: JitRL reaches a mean score of 53.0 with the same Qwen3-32B backbone, making the tuned GRPO result competitive.The comparison concerns the Zork1 validation setting.
  • Sample cost: 3,200 training episodes are required by the best GRPO configuration, over 64× more than JitRL’s 50 trajectories.The original GRPO setting still requires 400 episodes, or 8× more than JitRL.

R. Cost Analysis.

The cost analysis contrasts expensive gradient-based training with JitRL’s memory-based policy optimization and examines sensitivity, retrieval overhead, and GRPO resource requirements. The supplied passages emphasize that JitRL’s retrieval overhead remains negligible while GRPO can require many more trajectories.

  • Training cost: 154 hours of computation are required to train WebRL on a 16 × NVIDIA H200 GPU cluster.The process includes a 10-hour SFT phase followed by an eight-phase reinforcement-learning cycle.
  • Inference overhead: JitRL’s retrieval overhead remains negligible compared with LLM inference time as its memory bank grows.The evaluation measures retrieval latency and average score across memory bins on Library over 50 episodes.
  • JitRL procedure: JitRL retrieves similar experiences, estimates values and advantages, and applies an additive logit update before sampling actions.The algorithm combines LLM candidates with memory-derived actions and updates logits using normalized advantages.
  • JitRL procedure: The normalized advantage is computed as A(s, a) = Q(s, a) − V(s) before the closed-form policy update.The update modifies logits according to z′(s, a) ← z(s, a) + β · Ã(s, a).
Loading 2601.18510v3…