Context-Aware RL for Agentic and Multimodal LLMs

Abstract: LLMs often fail when answering requires identifying a small but decisive piece of evidence within a long or complex context, such as a single line in a tool trace or a subtle detail in an image. We propose ContextRL, a context-aware reinforcement learning (RL) method that improves long-horizon reasoning and multimodal performance through an \emph{indirect} auxiliary objective. Instead of supervising only the final answer, ContextRL presents the model with a query, an answer, and two highly similar contexts, and rewards it for selecting the context that supports the query--answer pair, thereby encouraging fine-grained grounding. We construct contrastive context data in two domains: for coding agents, trajectories serve as contexts, yielding 1k pairs built via condition filtering; for multimodal reasoning, images serve as contexts, yielding 7K pairs built via generative editing and similarity search. ContextRL achieves average gains of +2.2% over standard GRPO on 5 long-horizon benchmarks, and +1.8% across 12 diverse visual question answering benchmarks. To disentangle the effect of the proposed objective from that of additional data, we compare against data-augmentation baselines that repurpose the same contrastive contexts as standard query--context--answer examples. These baselines provide little to no improvement, showing that the gains arise from the proposed context-selection objective rather than from the contrastive data alone.

• The paper proposes ContextRL, which augments traditional RL with a context-selection auxiliary objective to improve evidence grounding in both coding and vision tasks.
• It demonstrates consistent improvements over standard RL baselines across multiple long-horizon reasoning benchmarks by leveraging contrastive pairs.
• Empirical results reveal that performance gains stem from the training objective design rather than data augmentation alone, ensuring robust context-driven predictions.

Context-Aware Reinforcement Learning for Agentic and Multimodal LLMs

Motivation and Problem Definition

LLMs and their multimodal variants (MLLMs) have exhibited substantial progress in long-horizon reasoning, agentic tool use, and integrated multimodal understanding. However, their performance often hinges on extracting decisive evidence from complex or lengthy contexts. Failures in grounding predictions to specific context details are prevalent, manifesting as context unawareness—errors arise not from lack of information, but from inadequate model attention or inference with respect to available evidence.

Figure 1: Context unawareness in action: agentic LLMs lose variable definitions across edits; multimodal LLMs misinterpret key visual evidence despite its presence.

The empirical gap between closed-source and state-of-the-art open-source models on a controlled “contrastive context selection” task reveals an unsettled challenge: strong LLMs can perform nearly at chance, even when benchmark metrics are competitive. This indicates that high downstream accuracy can mask serious deficits in context-sensitive grounding.

Methodology: ContextRL

The proposed method, Context-Aware Reinforcement Learning (ContextRL), augments standard RL post-training with an auxiliary context-selection objective. The central idea is to expose the model, for a given query–answer pair $(Q, A)$, to two nearly indistinguishable contexts: one that supports $A$ (“anchor”), and one that subtly contradicts it (“confounder”). The policy is rewarded for preferentially selecting the supporting context. This contrastive scheme targets evidence-level grounding within both text-only agentic traces and multimodal (image) inputs, all within the same RL framework.

Figure 2: ContextRL overview: the context-selection auxiliary objective integrates with GRPO, contrasting supporting and confounding contexts for each query–answer.

Contrastive pairs are constructed meticulously for both settings:

• Agentic (coding): Mining from 66k trajectories, only 1.5% survive aggressive filtering to ensure minimal surface-level confounds and maximal semantic challenge. Each pair is closely matched on repository, commit, modified file/function, and is exhaustively screened for shortcut cues.
• Multimodal (vision): Around 700 pairs are obtained via controlled generative editing and a further 6,300 via similarity-based retrieval on Qwen3-VL-Embedding, combining high perceptual similarity with label disparity.

  Figure 3: Construction pipeline for contrastive pairs in agentic and multimodal settings—stepwise filters and generative editing/retrieval ensure hard negative context formation.

The overall training objective is

$$\mathcal{L}(\theta) = \mathbb{E}_{x \sim \mathcal{D}_{\mathrm{RL}}}[\mathcal{L}_{\mathrm{GRPO}}(x;\theta)] + \lambda \, \mathbb{E}_{z \sim \mathcal{D}_{\mathrm{CA}}}[\mathcal{L}_{\mathrm{CA}}(z;\theta)],$$

where $\mathcal{L}_{\mathrm{CA}}$ is a logit-level contrastive loss explicitly encouraging robust context selection.

Empirical Findings

Long-Horizon Reasoning (Agentic Setting)

ContextRL was instantiated on Klear-AgentForge-8B and Qwen3-8B using 1k mined contrastive trajectory pairs. Across five agentic/long-context benchmarks, ContextRL demonstrates consistent improvements (+2.2% to +5.8% over RL baselines, outperforming even much larger models on SWE-Bench).

Notable empirical claim: Outcome-based RL (GRPO) sometimes regresses relative to the base model on out-of-distribution tasks, whereas ContextRL maintains or improves performance universally.

Multimodal Reasoning

Applied to Qwen2.5-VL-7B and Qwen3-VL-8B (with 7k contrastive image pairs in a 45k-example set), ContextRL delivers +2.0% to +1.6% average improvements across 12 benchmarks. The gains are distributed across perception-heavy, reasoning-heavy, and mixed tasks, and ContextRL consistently outperforms both standard RL and alternative perception-oriented RL like PAPO.

Figure 4: Only ContextRL achieves joint gains in context-selection accuracy and end-task performance; DA-SFT learns selection but fails at downstream reasoning; DA-RL fails at both.

Data Augmentation vs. Objective Design

Controls using the same contrastive data as either supervised fine-tuning (DA-SFT) or outcome-based RL (DA-RL) provide crucial negative results:

• DA-SFT: Collapses policy performance completely in agentic settings, despite excelling in context selection.
• DA-RL: Produces negligible gains, highlighting RL sparsity.
• ContextRL: Only the auxiliary loss formulation unlocks the benefit of contrastive data.

This analysis pinpoints that the primary contributor to performance gains is the training objective—i.e., contrastive context selection—not the data alone.

Ablation and Sensitivity Analysis

• Loss coefficient $\lambda$: Performance peaks at intermediate values; high $\lambda$ leads to over-prioritization of context selection and degrades main task performance.
• Proportion of contrastive data: An inverted-U relationship; 15% provides the optimal balance.
• Generation budget (response length): Increased context length typically aids, but excessive length can introduce noise on specific tasks.

  Figure 5: Accuracy as a function of contrastive data proportion; a moderate percentage elicits maximal benefit.

  

  Figure 6: Response length ablation; 4096 tokens provides the best overall trade-off across benchmarks.

  

  Figure 7: Excessive weighting of the context-awareness auxiliary loss ($\lambda$) hurts performance.

Dataset and Scalability Considerations

The multimodal training set integrates examples from diverse mathematical, scientific, and perceptual subdomains, ensuring that improvements are not domain-specific artifacts. The agentic contrastive set is precisely constructed for maximal semantic challenge and minimal artifact solvability.

Figure 8: Subgroup breakdown of the multimodal training set highlighting diversity in visual reasoning sources.

Theoretical and Practical Implications

Practical: ContextRL is architecturally lightweight and requires no additional reference models or large-scale manual annotation. Its applicability is demonstrated in both code-based agentic workflows and high-dimensional vision–language pipelines. The modality-agnostic nature suggests broad applicability in complex, information-rich domains, including tool-use agents and perceptually grounded systems.

Theoretical: This work decouples answer correctness from context-evidence association, showing that explicit process-level supervision—through context selection—addresses a fundamental limitation of outcome-based RL: the latter cannot reliably ensure that outputs are grounded in supplied evidence, even when end-task metrics are strong.

Comparison to related approaches—such as verifiable-context RL, process reward shaping, and contrastive SFT—clarifies that ContextRL introduces a distinct axis of supervision: context selection for a fixed $(Q, A)$, not answer selection for a fixed context.

Limitations and Future Directions

Current empirical validation covers only models up to the 8B–10B parameter scale, with primary focus on the Qwen family. Broader evaluation at ultra-large scales and across alternative architectures remains an open avenue. The contrastive mining process, while precise, is computationally intensive, especially for trajectory-based agentic domains.

Potential future developments include:

• Integration with architectural long-context enhancements and memory-augmented models.
• Automated mining of harder negative contexts at larger scales.
• Application to dialogue, instruction-following, and multi-agent scenarios where context-faithfulness is safety-critical.

Conclusion

ContextRL introduces a context-aware reinforcement learning procedure, coupling standard outcome-based RL with a context-selection auxiliary loss. Systematic experiments in both agentic (long-horizon, tool-use) and multimodal (vision–language) domains demonstrate that context-grounding can be substantially enhanced without additional architectural or data dependencies. The empirical and analytical results establish that the critical ingredient is not contrastive data per se, but an auxiliary objective engineered to directly supervise the link between answers and supporting context, thus strengthening the reliability and interpretability of LLM predictions across complex reasoning domains.

Reference: "Context-Aware RL for Agentic and Multimodal LLMs" (2606.17053)