JitRL: Just-in-Time Reinforcement Learning

July 14, 2026 · View on GitHub

Continual Learning in LLM Agents Without Gradient Updates

🎉 Accepted at ICML 2026 (Spotlight) • arXiv:2601.18510

Overview • Installation • Quick Start • Architecture • Results • Citation

Overview

While LLM agents excel at general tasks, they struggle with continual adaptation due to frozen weights after deployment. Conventional RL offers a solution but incurs prohibitive computational costs and the risk of catastrophic forgetting.

JitRL vs Standard RL

Standard RL performs gradient updates during training; JitRL operates at test time by retrieving relevant trajectories to estimate advantages and refine output logits.

JitRL is a training-free framework that enables test-time policy optimization without any gradient updates. Instead of updating parameters, JitRL:

Maintains a dynamic, non-parametric memory that stores experience trajectories as <state, action, reward> triplets
Retrieves relevant trajectories given the current state to estimate action advantages on-the-fly
Directly modulates the LLM's output logits based on these advantage estimates

We theoretically prove that this additive update rule is the exact closed-form solution to the KL-constrained policy optimization objective.

JitRL Framework

JitRL Framework: The inference stream retrieves experiences from memory to estimate advantages and adjust action logits; the memory update stream stores evaluated trajectories with step-wise returns.

Supported Environments

Environment	Domain	Tasks	Description
Jericho	Text Adventure Games	16+ games	Interactive fiction games (Zork, Library, etc.)
WebArena	Web Automation	812 tasks	Real-world web tasks (shopping, admin, maps, etc.)

Key Features

Jaccard Similarity Matching: N-gram based similarity for trajectory retrieval
LLM-based Step Scoring: Automatic evaluation of action quality
Dynamic Prompt Generation: Adaptive prompts based on task and history
Multi-model Support: OpenAI, Anthropic, Google, and open-source models via OpenRouter

Installation

Prerequisites

Python 3.10+
OpenRouter API key (for LLM inference)

Clone Repository

git clone https://github.com/your-username/JitRL.git
cd JitRL

Jericho Setup (Text Adventure Games)

cd Jericho

# Install dependencies
pip install jericho openai tiktoken numpy python-dotenv

# Download game ROMs (required)
# Place .z5/.z8 files in jericho-games/ directory
# Games available: zork1, zork3, library, detective, etc.

WebArena Setup (Web Automation)

cd WebArena

# Install dependencies
pip install browsergym-core browsergym-experiments
pip install openai tiktoken numpy langchain pillow

Environment Configuration

Create a .env file in the project root:

# OpenRouter API (for LLM inference)
OPENROUTER_API_KEY=your_openrouter_key

# WebArena URLs (if running locally)
WA_SHOPPING=http://localhost:7770
WA_SHOPPING_ADMIN=http://localhost:7780/admin
WA_REDDIT=http://localhost:9999
WA_GITLAB=http://localhost:8023
WA_MAP=http://localhost:3000
WA_WIKIPEDIA=http://localhost:8888

Quick Start

Jericho: Text Adventure Games

cd Jericho

# Run JitRL agent on Zork 1 for 10 episodes
python main.py --game_name zork1 --agent_type jitrl --eval_runs 10

# Run with cross-episode memory disabled (baseline)
python main.py --game_name zork1 --agent_type jitrl --no-enable_cross_mem

# Use different LLM model
python main.py --game_name zork1 --llm_model openai/gpt-4o --eval_runs 10

Key Arguments:

Argument	Default	Description
`--game_name`	`library`	Game to play
`--agent_type`	`jitrl`	Agent: `jitrl`, `naive`, `awm`
`--eval_runs`	`50`	Number of episodes
`--llm_model`	`google/gemini-2.5-flash`	LLM model
`--env_step_limit`	`50`	Max steps per episode
`--enable_cross_mem`	`True`	Enable cross-episode memory
`--gamma`	`0.5`	Discount factor for rewards

Note on reproducibility: The Gemini-2.5-Flash results in the paper were obtained with the google/gemini-2.5-flash-preview-09-2025 snapshot on OpenRouter, which has since been retired. The defaults now point to the stable google/gemini-2.5-flash; reproduced numbers may differ slightly due to the model version change.

WebArena: Web Automation

cd WebArena

# Run single task with memory agent
python test_webarena_lite.py --tasks 0 --repeat 5 --model gpt-4o

# Run multiple tasks in parallel
python test_webarena_lite.py --start 0 --end 10 --workers 4 --repeat 3

# Run with screenshots for vision-capable models
python test_webarena_lite.py --tasks 68 --model gpt-4o \
    --use_screenshot_action --use_screenshot_eval

# Run without memory (baseline)
python test_webarena_lite.py --tasks 0 --repeat 10 --disable_memory

Key Arguments:

Argument	Default	Description
`--tasks`	-	Comma-separated task IDs (e.g., `0,1,2`)
`--start/--end`	-	Task ID range
`--repeat`	`1`	Episodes per task
`--workers`	`1`	Parallel workers
`--model`	`gpt-4o`	LLM model
`--max_steps`	`30`	Max steps per episode
`--disable_memory`	`False`	Disable memory system

Architecture

Project Structure

JitRL/
├── Jericho/                          # Text Adventure Games
│   ├── main.py                       # Entry point
│   ├── console_play.py               # Interactive play mode
│   └── src/
│       ├── jitrl_agent.py            # JitRL agent: retrieval + advantage + logit modulation (850 lines)
│       ├── naive_agent.py            # Simple baseline agent
│       ├── awm_agent.py              # Agent Workflow Memory
│       ├── cross_episode_memory.py   # Core memory system (900 lines)
│       ├── evaluation.py             # Game evaluator
│       ├── env.py                    # Jericho environment wrapper
│       ├── utils.py                  # LLM utilities (1150 lines)
│       ├── openai_helpers.py         # API helpers
│       └── prompt_update_with_history.py  # Prompt optimization
│
├── WebArena/                         # Web Automation
│   ├── test_webarena_lite.py         # Main test script
│   ├── run.py                        # Legacy entry point
│   ├── memory_agents/
│   │   ├── jitrl_agent.py            # JitRL BrowserGym agent (1450 lines)
│   │   ├── dynamic_prompting.py      # Prompt generation (760 lines)
│   │   └── utils/
│   │       ├── cross_episode_memory.py  # Memory system (1200 lines)
│   │       ├── openai_helpers.py     # LLM utilities
│   │       ├── llm_utils.py          # Parsing utilities
│   │       └── utils.py              # General utilities (2100 lines)
│   ├── autoeval/                     # Evaluation system
│   │   ├── evaluator.py              # LLM-based evaluator
│   │   ├── enhanced_evaluator.py     # Rule-based evaluator
│   │   └── live_evaluator.py         # Real-time evaluation
│   ├── config_files/                 # Full WebArena (812 tasks)
│   └── config_files_lite/            # WebArena-Lite (165 tasks)
│
└── README.md

Supported Models

JitRL supports any model available through OpenRouter:

Provider	Models	Notes
OpenAI	`gpt-4o`, `gpt-4-turbo`, `gpt-4o-mini`	Best for structured output
Anthropic	`claude-3.5-sonnet`, `claude-3-opus`	Strong reasoning
Google	`gemini-2.5-flash`, `gemini-2.5-pro`	Cost-effective
Meta	`llama-3.1-70b`, `llama-3.1-405b`	Open-source
Mistral	`mistral-large`, `mixtral-8x22b`	Fast inference

Results

Jericho Text Adventures

Game	Naive Agent	Memory Agent	Improvement
Zork1	35.2	52.8	+50%
Library	18.5	28.3	+53%
Detective	180	265	+47%

Scores averaged over 50 episodes

WebArena-Lite

Model	Without Memory	With Memory	Improvement
GPT-4o	32.5%	41.2%	+27%
Claude-3.5	28.8%	38.5%	+34%
Gemini-2.5	25.3%	34.7%	+37%

Success rate on 165 WebArena-Lite tasks

Advanced Usage

Clearing Memory

from src.cross_episode_memory import CrossEpisodeMemory

memory = CrossEpisodeMemory("output/game/agent/model")
memory.clear_memory(save_to_disk=True)

Custom Retrieval Parameters

# Adjust retrieval sensitivity
python main.py --game_name zork1 \
    --retrieval_top_k 10 \
    --retrieval_threshold 0.8 \
    --gamma 0.95

Prompt Optimization

# Enable automatic prompt updates
python main.py --game_name zork1 \
    --update_guiding_prompt \
    --use_history_prompt

Troubleshooting

Common Issues

OpenRouter rate limits: Use exponential backoff (built-in)
Memory errors: Reduce --retrieval_top_k or --max_memory
WebArena connection: Ensure Docker containers are running

Logging

Logs are saved to:

Jericho: output/{game}/{agent}/{model}/{timestamp}/
WebArena: results/{task_id}/ or --log_dir

Citation

If you use JitRL in your research, please cite:

@article{li2026just,
  title={Just-In-Time Reinforcement Learning: Continual Learning in LLM Agents Without Gradient Updates},
  author={Li, Yibo and Lin, Zijie and Deng, Ailin and Zhang, Xuan and He, Yufei and Ji, Shuo and Cao, Tri and Hooi, Bryan},
  journal={arXiv preprint arXiv:2601.18510},
  year={2026}
}

Acknowledgments

Jericho - Text adventure game framework by Microsoft Research
BrowserGym - Web automation framework by ServiceNow
WebArena - Web agent benchmark by CMU

License

This project is for research purposes only. See LICENSE for details.