💬 NLP

Context-Aware RL for Agentic and Multimodal LLMs

The paper proposes ContextRL, a context-aware reinforcement learning method that improves long-horizon reasoning and multimodal performance in LLMs by training models to select the correct context from highly similar pairs via an indirect auxiliary objective, rather than relying solely on final answer supervision.

Original authors: Peiyang Xu, Bangzheng Li, Sijia Liu, Karthik R. Narasimhan, Pramod Viswanath, Prateek Mittal, Xingyu Fu

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Peiyang Xu, Bangzheng Li, Sijia Liu, Karthik R. Narasimhan, Pramod Viswanath, Prateek Mittal, Xingyu Fu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery. You have a massive stack of evidence files (the "context") and a specific question to answer. The problem is, most modern AI detectives are great at guessing the answer, but they often fail to point to the exact page in the file stack that proves they are right. They might get the answer right by luck or by remembering a pattern, but they can't show you the "smoking gun" in the evidence.

This paper, titled "Context-Aware RL for Agentic and Multimodal LLMs," introduces a new training method called CONTEXT-RL to fix this. It teaches AI models not just what to answer, but where to look in the evidence to justify that answer.

Here is a simple breakdown of how it works, using everyday analogies:

1. The Problem: The "Guessing" Detective

The authors noticed that even very smart AI models often suffer from "context unawareness."

In Coding: Imagine an AI is fixing a computer program. It sees a long list of instructions (the context). It decides to delete a variable named i. But if it looked closely at the context, it would see that i is used later in the code. Because it didn't "ground" its decision in the specific line of code, it breaks the program.
In Images: Imagine an AI looking at a graph. It needs to read a specific number. It might guess "2" because that's a common number, even though the graph clearly shows "3." It missed the tiny visual detail that was right in front of it.

The paper tested this by showing models a question, an answer, and two very similar stories (contexts). One story supported the answer, and the other was a "fake" story that looked almost the same but didn't support the answer. Surprisingly, many smart open-source models couldn't tell the difference and picked the wrong story almost as often as if they were guessing randomly.

2. The Solution: The "Spot the Difference" Game

To fix this, the authors created a new training game called CONTEXT-RL.

Instead of just telling the AI, "You got the answer right, here is a reward," they added a second, stricter rule: "You must also point to the correct evidence file."

The Setup: The AI is given a question and an answer. Then, it is shown two nearly identical "worlds" (contexts).
- World A: Contains the specific clue that proves the answer is correct.
- World B: Looks almost the same, but the clue is missing or changed, making the answer wrong.
The Goal: The AI gets a bonus reward not just for solving the problem, but for correctly identifying World A as the one that supports the answer.

It's like a teacher giving a student a math problem. A normal teacher says, "Good job, you got 5." The CONTEXT-RL teacher says, "Good job, but also show me the exact line in the textbook where you found the formula. If you can't point to it, you didn't really understand it."

3. How They Built the Training Data

To teach this game, they had to create thousands of "Spot the Difference" puzzles:

For Coding Agents: They took real-world coding tasks and found pairs of coding histories that were almost identical, except for one tiny, crucial difference in the code. They masked the actual code changes so the AI couldn't cheat by just reading the final fix; it had to understand the process.
For Images: They used AI tools to edit photos. For example, they took a picture of a chart and slightly changed one number so the answer to a question would flip from "Yes" to "No." They made sure the rest of the picture looked exactly the same, forcing the AI to look at that specific detail.

4. The Results: Why It Matters

The paper tested this method on two types of tasks:

Long-Horizon Coding: Agents that have to write code over many steps.
Multimodal Reasoning: Agents that have to look at images and answer questions.

The findings were clear:

Better Performance: Models trained with CONTEXT-RL got significantly better scores on difficult benchmarks than models trained with standard methods.
The Secret Sauce: The authors proved that the improvement didn't come just from having more data. They tried feeding the same "Spot the Difference" data to models using standard training methods, and it didn't work (or even made things worse). The magic was in the specific way they trained the model to select the right context.

The Takeaway

Think of standard AI training as teaching a student to memorize the final answer key. CONTEXT-RL teaches the student to be a detective who knows how to find the evidence in a messy room. It doesn't just want the right answer; it demands that the answer be grounded in the specific details provided, making the AI more reliable and less likely to make mistakes when the context is complex or subtle.

Technical Summary: Context-Aware RL for Agentic and Multimodal LLMs

Problem Statement

Large Language Models (LLMs), particularly those operating in agentic (tool-using) and multimodal (vision-language) settings, frequently fail when tasks require identifying a small but decisive piece of evidence within a long or complex context. This failure mode, termed context unawareness, occurs when a model produces a locally plausible answer that is not grounded in the specific evidence provided (e.g., modifying a variable defined elsewhere in a code trajectory, or misreading a specific value in a dense chart).

While modern models often achieve high scores on standard benchmarks, the authors introduce a contrastive context probe to reveal a significant gap in context grounding. In this probe, a model is given a query, a candidate answer, and two highly similar contexts (one supporting the answer, one supporting a different answer). The model must select the correct context. Results show that while proprietary models (e.g., GPT-5.4, Claude Opus 4.7) perform reliably, strong open-source models (e.g., Qwen3-VL, Qwen3.5) often perform close to random chance, despite competitive performance on standard benchmarks. This suggests that standard benchmark metrics obscure failures in grounding decisions in sparse contextual evidence.

Methodology: CONTEXTRL

The authors propose CONTEXTRL, a post-training framework that augments standard Reinforcement Learning (RL) with an explicit, auxiliary context-selection objective. Unlike standard RL, which optimizes only for the correctness of the final output (outcome rewards), CONTEXTRL trains the model to identify which context supports a given answer.

1. Contrastive Data Construction

The method relies on constructing contrastive context pairs $(C^+, C^-)$ associated with a query $Q$ and an anchor answer $A$ :

Agentic Setting (Coding): Trajectories serve as contexts. The authors mine 1,000 high-quality pairs from 66,000 SWE-smith trajectories. Pairs are constructed by filtering for trajectories from the same repository, commit, file, and target function, differing only in a small, decisive code region. Edit commands are masked to prevent leakage.
Multimodal Setting (Vision): Images serve as contexts. The authors construct 7,000 pairs using two strategies:
- Generative Editing: For natural images, a localized edit is applied to change the answer while preserving the global scene structure.
- Similarity-Based Retrieval: For structured images (charts, diagrams), highly similar images supporting different answers are retrieved using embedding models.
- Rigorous verification (using GPT-5.4 and manual inspection) ensures pairs are free of superficial shortcut cues or artifacts.

2. Training Objective

CONTEXTRL combines the standard Group Relative Policy Optimization (GRPO) objective with a Context-Awareness Loss ( $L_{CA}$ ).

Standard GRPO ( $L_{GRPO}$ ): Optimizes for correct task outcomes (e.g., passing test cases in coding, exact match in VQA).
Context-Awareness Loss ( $L_{CA}$ ): For a contrastive instance $(Q, A, C^+, C^-)$ , the model is presented with the query, answer, and both contexts as options. The loss encourages the model to assign a higher logit to the token corresponding to the supporting context ( $C^+$ ) over the confounding context ( $C^-$ ).
Joint Optimization: The total loss is $L(\theta) = E[L_{GRPO}] + \lambda E[L_{CA}]$ , where $\lambda$ balances task performance and context awareness. The $L_{CA}$ is implemented as a margin-based loss with clipping to prevent large gradients from dominating training.

Key Contributions

Diagnostic Probe: The introduction of a controlled contrastive context probe that exposes a ~40-point gap in context grounding between proprietary and open-source models, revealing that standard benchmarks may not capture context-awareness failures.
CONTEXTRL Framework: A lightweight, modality-agnostic RL post-training method that adds a context-selection signal to outcome-based RL. It does not require architectural changes or large-scale human annotation.
Data Construction Pipelines: Novel pipelines for generating high-quality contrastive data in both agentic (trajectory filtering) and multimodal (generative editing and retrieval) domains, with rigorous artifact rejection.
Disentanglement of Data vs. Objective: A critical ablation study demonstrating that the performance gains arise from the training objective (context selection) rather than the mere presence of contrastive data. Standard data augmentation baselines (SFT or outcome-based RL on the same data) fail to improve performance or cause catastrophic policy collapse.

Experimental Results

The method was evaluated across 17 benchmarks in two domains:

Long-Horizon Agentic Coding

Benchmarks: SWE-Bench Verified/Lite (In-Distribution), LiveCodeBench v6, LongBench v2, and Needle-in-a-Haystack (Out-of-Distribution).
Results: CONTEXTRL achieved average gains of +2.2% over standard GRPO across five benchmarks. Notably, it demonstrated strong generalization to Out-of-Distribution (OOD) tasks where standard RL often regressed. For instance, on the Needle-in-a-Haystack task, standard GRPO performed worse than the base model, while CONTEXTRL improved upon it.
Scale Efficiency: When trained on Klear-AgentForge-8B, CONTEXTRL outperformed the much larger Qwen3-32B and code-specialized Qwen3-Coder-30B on SWE-Bench.

Multimodal Reasoning

Benchmarks: 12 diverse benchmarks covering mathematical reasoning, general understanding, fine-grained perception, scientific reasoning, and real-world scene understanding.
Results: CONTEXTRL achieved average gains of +1.8% over standard GRPO across all 12 benchmarks on both Qwen2.5-VL-7B and Qwen3-VL-8B. Improvements were consistent across all sub-categories (perception-heavy, reasoning-heavy, and mixed), indicating that enhanced visual grounding propagates to higher-level reasoning.

Ablation: Data Augmentation vs. Objective

The authors compared CONTEXTRL against two baselines using the same contrastive data:

DA-SFT (Supervised Fine-Tuning): Learning to select the context via SFT caused catastrophic policy collapse in the agentic setting (resolve rates dropped to 0% or <2%) and negligible gains in multimodal settings. This suggests SFT on short selection tasks disrupts the long-horizon interaction policy.
DA-RL (Outcome-based RL): Treating context selection as a binary reward task yielded little to no improvement over the standard RL baseline.
Conclusion: The gains are specific to the context-selection objective integrated within the on-policy GRPO loop, which provides dense, bounded gradients without disrupting the primary task policy.

Significance and Claims

The paper claims that failures in long-horizon reasoning and visual understanding often stem from a common limitation: models do not reliably ground predictions in available evidence. CONTEXTRL addresses this by decoupling "what answer to produce" from "which context supports it."

The authors position context selection as a simple, lightweight, and modality-agnostic mechanism for improving context grounding. The work emphasizes that the method of integrating the signal is the key factor; the contrastive data alone is insufficient and can even be detrimental if used with standard augmentation techniques. The results suggest that a single auxiliary objective can significantly improve the reliability of agentic and multimodal systems without requiring massive scale or architectural overhaul.