Source-linked AI summary

Context-Aware RL for Agentic and Multimodal LLMs

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

arXiv:2606.17053v1cs.CLcs.CV

TL;DR

LLMs can overlook sparse but decisive evidence in long trajectories and dense images, producing predictions that are not grounded in available context. ContextRL adds a context-selection objective to reinforcement learning and consistently improves long-horizon reasoning and multimodal understanding over standard GRPO across 17 benchmarks.

  • Problem

    LLMs often fail to ground decisions in sparse but decisive evidence available within long trajectories, tool traces, and dense images.

  • Method

    ContextRL augments reinforcement-learning post-training with an objective rewarding selection of the context that supports a given query–answer pair.

  • Results

    Across 5 long-horizon and 12 multimodal benchmarks, CONTEXTRL consistently improves over GRPO, including +2.0% on average for Qwen2.5-VL.

  • Takeaways & Limitations

    The results support context selection as a lightweight auxiliary signal whose benefits arise from the training objective rather than contrastive data alone.

  • Takeaways & Limitations

    Experiments use models under 10B parameters, mostly from the Qwen family, so generality to substantially larger models and other model families remains unvalidated.

Abstract

from arXiv · show

Large language models (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.

1 Introduction

ContextRL addresses a shared failure in agentic and multimodal systems: overlooking sparse but decisive evidence in long trajectories or dense images. It adds a context-selection objective to reinforcement learning, producing improvements across long-horizon reasoning and multimodal benchmarks beyond contrastive data alone.

  • Problem: Models can make context-inconsistent decisions when they overlook sparse but decisive evidence in long agent trajectories or dense images.Examples include missing an early observation, a tool-trace line, or a subtle visual detail.
  • Diagnostic probe: A controlled probe presents a query, candidate answer, and two closely matched contexts, revealing a roughly 40-point gap between proprietary and open-source models.Qwen-3 (VL) 8B and Qwen-3.5 9B perform close to random choice despite competitive standard-benchmark performance.
  • Method: ContextRL augments standard reinforcement learning with an auxiliary reward for selecting the context that supports a given answer.Contrastive context pairs are constructed through condition search and generative editing.
  • Results: +3.2% and +1.5% are ContextRL’s average improvements over GRPO across five agentic and long-context benchmarks using Klear-AgentForge-8B and Qwen3-8B, respectively.The long-horizon experiments use 1k contrastive trajectory pairs constructed from SWE-smith trajectories through condition filtering.
  • Findings: The improvements show that a single context-selection objective benefits both long-horizon reasoning and multimodal understanding without architectural changes or large-scale human annotation.Comparisons with data-augmentation baselines indicate that the gains arise from the objective rather than contrastive data alone.

2 CONTEXTRL

ContextRL augments standard RL with a context-selection signal that trains models to identify which of two similar contexts supports a query–answer pair. It constructs contrastive trajectory and image data, then combines task rewards with a modality-agnostic context-awareness loss.

  • Core framework: ContextRL rewards models for selecting the context that supports a given answer, complementing standard RL rewards for correct outputs.The framework uses contrastive pairs (C+, C−) across agentic and multimodal settings.
  • Agentic contrastive data: Agentic pairs comprise queries, reference patches, and reasoning trajectories whose decisive evidence must be identified within tool interactions and observations.Pairs are mined from trajectories and filtered to share repository, commit, file, and target function or class while representing distinct related issues.
  • Multimodal contrastive data: 7k contrast image pairs combine generative editing for natural images with similarity-based retrieval for structured or text-rich images.Retrieval uses a visual similarity threshold of αI = 0.85, and filtering yields 6,300 high-quality pairs from over 200k candidates.
  • Joint objective: The joint objective combines GRPO task optimization with context alignment through a balance parameter λ > 0, using the same formulation for agents and multimodal inputs.GRPO uses test-case success for agentic tasks and exact answer match for multimodal tasks, while the auxiliary loss enforces grounding.
  • Context-awareness objective: The context-awareness loss presents Q, A, and randomized positive and negative contexts as a two-way multiple-choice prompt, encouraging higher confidence in C+.It compares teacher-forced next-token logits and uses sigmoid margin clipping controlled by c > 0.

3 Long Horizon Experiments

ContextRL is evaluated with two agentic coding base models across five long-horizon benchmarks, using standard GRPO and no-RL configurations as baselines. It consistently improves over outcome-based RL, generalizes to out-of-distribution long-context tasks, and is competitive with much larger reference models.

  • Experimental setup: Experiments use Qwen3-8B and Klear-AgentForge-8B, spanning a general-purpose and an agentic-coding-specialized base model.The training configurations compare Base, standard GRPO, and ContextRL; the training set contains 8k instances, including 7k standard coding tasks and 1k contrastive trajectory pairs.
  • Evaluation benchmarks: Evaluation covers five benchmarks: two in-distribution agentic-coding tasks and three out-of-distribution tasks spanning competitive programming, long-context QA, and targeted retrieval.The in-distribution benchmarks are SWE-Bench Verified and SWE-Bench Lite; the out-of-distribution benchmarks are LiveCodeBench v6, LongBench v2, and NIAH.
  • Main results: ContextRL improves over the RL baseline on every benchmark for both base models, demonstrating a robust training signal beyond outcome-only RL.The reported results cover all five long-horizon benchmarks and both base models.
  • Scale comparison: When trained from Klear-AgentForge-8B, ContextRL substantially outperforms Qwen3-32B, which is 4× larger, and code-specialized Qwen3-Coder-30B on SWE-Bench.This comparison indicates that the context-aware objective combined with a domain-adapted base can compensate for differences in scale and pretraining specialization.
  • Generalization: On every out-of-distribution benchmark and both base models, ContextRL outperforms the base model and RL baseline, with particularly pronounced effects on LongBench v2 and NIAH.On NIAH, standard GRPO regresses relative to the base model while ContextRL surpasses it; on LongBench v2, ContextRL improves both the overall and long subsets.
  • Ablations: Performance peaks at λ = 0.005, balancing context-awareness and task optimization.The coefficient ablation examines the effect of the context-awareness loss weight λ.

4 Multimodal Experiments

ContextRL improves multimodal performance through stronger visual grounding across two vision-language base models and 12 diverse benchmarks. Its gains consistently surpass standard GRPO across task categories, while small auxiliary-loss weights perform best.

  • Experimental setup: The study evaluates Qwen2.5-VL-7B-Instruct and Qwen3-VL-8B-Instruct as comparable-scale vision-language base models.
  • Experimental setup: The comparison isolates the learning objective among an untrained base model, standard GRPO, and GRPO augmented with ContextRL’s context-awareness loss.The multimodal training set comprises 45k examples, including 38k standard single-image task instances.
  • Evaluation benchmarks: Evaluation spans 12 benchmarks covering mathematical, general, fine-grained perceptual, scientific, and real-world scene understanding.The benchmarks include MathVista, MathVerse, MathVision, MMMU-Pro, MMMU, V*, MMStar, BLINK, ScienceQA, PhyX, OlympiadBench Physics, and MME-RealWorld Lite.
  • Main results: ContextRL outperforms standard GRPO on every benchmark for both base models and across all five sub-task categories.Improvements span perception-heavy, reasoning-heavy, and mixed tasks without category trade-offs.
  • Main results: The consistent gains across categories and model families suggest improved context grounding rather than category- or model-specific tuning.
  • Ablations: Optimal multimodal performance occurs at small auxiliary-loss weights of 0.001 or 0.005, while larger values degrade performance.At larger weights, the auxiliary context-awareness loss overwhelms the primary GRPO objective.

5 Comparison between Data Augmentation and CONTEXTRL

Standard data augmentation with the same contrastive data is ineffective or destabilizing, whereas CONTEXTRL’s objective extracts the useful signal without sacrificing task performance. Mechanism analyses show that high context-selection accuracy alone is insufficient; the signal must be integrated without disrupting the original policy.

  • Baselines: DA-SFT applies supervised context prediction before standard GRPO, while DA-RL mixes binary context-selection examples directly into the RL stream.Both methods use the same contrastive data as CONTEXTRL but incorporate it through standard augmentation objectives.
  • Agentic results: On agentic benchmarks, DA-SFT causes catastrophic policy collapse, reducing Klear-AgentForge-8B resolve rates from 28.0 / 21.7 to 6.4 / 1.3 and Qwen3-8B to 0.00 / 0.00.The collapse occurs despite the model learning the contrastive selection task.
  • Multimodal results: On multimodal benchmarks, DA-SFT averages 51.5 and 63.9 versus 51.4 and 64.1 for the RL baseline, while DA-RL yields only marginal gains of +0.4 / +0.2.These results indicate that both augmentation strategies are largely ineffective across the 12 benchmarks.
  • Multimodal results: CONTEXTRL improves multimodal average performance by +2.0 points on Qwen2.5-VL and +1.6 points on Qwen3-VL, improving every benchmark individually.The contrastive data alone is insufficient; the gains arise from how CONTEXTRL incorporates its signal during training.
  • Mechanism study: DA-SFT and CONTEXTRL raise selection accuracy to 85–93%, but only CONTEXTRL consistently improves downstream performance.This shows that context-selection skill is necessary but not sufficient, because it must be acquired without disrupting task-policy behaviors.
  • Mechanism study: CONTEXTRL couples high context awareness with task gains by using a bounded auxiliary objective inside on-policy GRPO, while its dense relative-preference signal avoids DA-RL’s sparse 0, 1 rewards.GRPO clipping, KL regularization, and the clipped margin objective constrain updates and preserve the original policy distribution.

6 Related Work

Prior work applies reinforcement-learning post-training to mathematical, coding, agentic, and multimodal settings, while separate approaches improve faithful context utilization through long-context retrieval, context rewards, and evidence-preserving compression. ContextRL builds on these lines of research.

  • RL Post-training in Agentic and Multimodal Settings: RL post-training has advanced mathematical and coding reasoning through methods including RLHF, DPO, GRPO, and DAPO.The passage identifies GRPO and DAPO as recent verifiable-reward methods driving progress.
  • RL Post-training in Agentic and Multimodal Settings: Recent work extends RL post-training to software-engineering agents through SWE-RL and DeepSWE.These methods explore RL as a training paradigm for agentic software engineering.
  • Context Utilization and Contrastive Supervision: LLMs often fail to use supplied context faithfully, motivating research on context utilization and contrastive supervision.The passage frames this as a growing body of work addressing failures in faithful context use.
  • Context Utilization and Contrastive Supervision: Long-context approaches include FILM’s information-intensive retrieval supervision, LongRLVR’s context rewards, and MemOCR’s memory- and layout-aware compression.These methods target long-context settings while preserving or rewarding access to relevant evidence.

7 Conclusion … A.2 Multimodal Data Sources

ContextRL adds a lightweight context-selection objective to outcome-based post-training and improves performance over GRPO across long-horizon and multimodal settings. Its contrastive data pipelines mine tightly matched agentic trajectories and multimodal context pairs, with gains attributed to the objective rather than contrastive data alone.

  • 7 Conclusion: ContextRL improves over GRPO in both long-horizon and multimodal settings across 17 benchmarks.Comparisons with data-augmentation baselines attribute the gains to the training objective rather than contrastive data alone.
  • 7 Conclusion: Context selection serves as a simple, consistently beneficial auxiliary signal for improving context grounding.The framework adds this objective to outcome-based post-training.
  • A.1 Detailed Procedure for Mining Agentic Contrast Trajectories: Contrastive trajectory pairs are mined from 66k SWE-smith trajectories through repository, commit, file, and target-function or class filtering.The filters ensure shared code bases and file-level comparability while preventing direct edit-command shortcuts.
  • A.1 Detailed Procedure for Mining Agentic Contrast Trajectories: The agentic selection task retains file views, test outputs, error messages, and reasoning traces as valid context rather than masking them.Filtering is designed to prevent solving the task by directly reading the edit command.
  • A.1 Detailed Procedure for Mining Agentic Contrast Trajectories: Only 1k pairs survive conservative filtering, GPT 5.4 verification, and manual inspection from the original 66k trajectories.Ambiguous pairs and examples with residual leakage are removed.
  • A.2 Multimodal Data Sources: Multimodal contrast context pairs span five visual domains, using generative editing for natural images and similarity-based retrieval for structured visual inputs.The source datasets cover chart, geometry, non-geometric math, science, and natural-image domains.
  • A.2 Multimodal Data Sources: The multimodal sources include ChartQA and DVQA; Geo170K, Geometry3K, GeomVerse, GeoQA3, MAVIS-Geometry, and R-CoT; and multiple math, science, and natural-image datasets.The listed sources include ICON-QA, K12, MAVIS-Function, MMK12, MM-Math, AI2D, M3CoT, ScienceQA, and Visual CoT.

A.3 Dataset Difficulty and Quality Control: Screening for Artifacts and Shortcut Cues

The dataset is designed to discourage artifact- or shortcut-based context selection through explicit verifier filters, near-identical hard negatives, and aggressive rejection. Its validity is further supported by gains on out-of-distribution benchmarks containing none of the constructed inputs, where artifact detection would be inapplicable.

  • Artifact and shortcut filtering: Verifier criteria reject trajectory pairs with formatting, length, token-leakage, inspection, or masking disparities that could provide shortcut cues.The trajectory verifier rejects pairs with large length or formatting disparities, patch-specific tokens appearing in only one trajectory, asymmetric inspection, or inconsistent masking.
  • Hard-pair construction: Hard negatives share the same repository, commit, modified file, and target, differing only in a small decisive code region.Edit commands are masked by <PATCH_MASKED>, making paired agentic contexts nearly identical at the token level.
  • Aggressive rejection: Each pipeline retains only a small, high-precision fraction of candidates after automatic verification and manual review.The filtering funnel is summarized in Table 5.
  • Evidence against artifact exploitation: ContextRL improves on OOD benchmarks containing none of the constructed inputs, including 12 multimodal benchmarks using natural, unedited single images.The agentic OOD suite includes LiveCodeBench, LongBench v2, and NIAH; the multimodal benchmarks use the standard non-selection format.

B Dataset Composition · C Prompt Templates

The paper uses a 45,000-example multimodal training set split between standard GRPO examples and contrastive context-selection instances, alongside task-specific prompts for trajectory, image, verification, coding, and multimodal reasoning. These templates define the inputs, selection criteria, validity checks, and required output formats for each setting.

  • B Dataset Composition: The multimodal training set contains 45,000 examples: 38,000 standard single-image GRPO examples and 7,000 two-image contrastive context-selection instances.The split is 85% standard task examples and 15% contrastive examples.
  • B Dataset Composition: The dataset spans Geometry, Non-geometric Math, Chart, Science, and Natural Image subgroups, with per-source counts reported within each subgroup.Figure 6 shows subgroup composition, while Table 6 provides the source-level breakdown.
  • C.1 Trajectory Selection Prompt: Trajectory selection presents a code patch and two masked candidate trajectories, requiring the model to identify which trajectory produced the patch and output only A or B.The prompt asks the model to compare each trajectory’s diagnosed bug with the patch’s fix.
  • C.2 Image Selection Prompt: Image selection places two images before a question–ground-truth-answer pair and asks the model to select the first or second image consistent with that reference.The prompt exposes the question and ground-truth answer as verification data and offers two image choices.
  • C.4 Contrast Image Pairs Verification Prompt: GPT 5.4 verifies contrast image pairs for visual coherence, correctness of the new answer, and preservation of question-irrelevant content, returning a validity decision and confidence score.The edited image must remain natural, answer A' must be unambiguous for Q, and unrelated regions must remain essentially unchanged.
  • C.5 Agentic Coding Prompt: The agentic coding prompt wraps a PR description and task instructions around an interactive shell workflow requiring reasoning plus exactly one bash command per response.The workflow directs agents to inspect, reproduce, edit, verify, and test changes while modifying regular source files rather than tests or configuration.
  • C.6 Multimodal Reasoning Prompt: The multimodal reasoning prompt supplies a question and options, requiring step-by-step thinking inside <think></think> and the final answer inside <answer></answer>.Curly-braced fields in templates are filled with per-instance content, while other text remains fixed.

D Training Details and Hyperparameters · D.1 Agentic Setting · D.2 Multimodal Setting

The paper uses framework-specific GRPO training with distinct agentic and multimodal setups. Agentic training uses mini-SWE-agent and 8k instances, while multimodal training uses 45k examples with contrastive image pairs matched against an equal-size RL baseline.

  • D.1 Agentic Setting: Agentic training uses GRPO in SkyRL with mini-SWE-agent for both sandboxed coding rollouts and evaluation.SkyRL is designed for long-horizon tool-use rollouts.
  • D.1 Agentic Setting: Table 7 reports the key training hyperparameters for the agentic coding setting.The supplied passage identifies the table but does not provide its parameter values.
  • D.2 Multimodal Setting: Multimodal training uses GRPO in Easy-R1 with 45k examples: 38k standard single-image QA examples for LGRPO and 7k contrastive image pairs for LCA.The multimodal setup uses the stated 45k-example composition.
  • D.2 Multimodal Setting: The multimodal RL baseline also trains on 45k standard QA examples for the same number of steps, replacing contrastive pairs with additional source-pool samples.Its 45k examples comprise the same 38k standard examples plus 7k additional samples.
  • D.2 Multimodal Setting: Table 8 reports the key training hyperparameters for the multimodal setting.The supplied passage identifies the table but does not provide its parameter values.

E Compute Resources · F Ablations · F.1 Agentic Setting

The paper uses a single 4× H200 node for both RL settings and reports additional compute for ablations and data construction. Ablations show agentic performance is sensitive to the context-awareness loss weight, peaking at λ = 0.005.

  • E Compute Resources: E Compute Resources: RL training uses one node with 4× NVIDIA H200 GPUs, NVLink, and at least 500 GB host RAM.The same configuration serves both multimodal and agentic experiments.
  • E Compute Resources: E Compute Resources: Each per-experiment cost covers one end-to-end RL run for one base model under one configuration.Ablations and data-augmentation comparisons each require an additional run of comparable cost.
  • E Compute Resources: E Compute Resources: Nano Banana 2 was queried approximately 10k times to generatively edit natural images.These queries supported multimodal data construction.
  • E Compute Resources: E Compute Resources: GPT-5.4 was queried approximately 10k times to automatically verify edited images.The verifier was used during data construction.
  • F Ablations: F Ablations: The paper provides detailed ablation studies for both agentic and multimodal settings.The reported ablations include the agentic context-awareness loss sweep.
  • F.1 Agentic Setting: F.1 Agentic Setting: λ = 0.005 produces the best performance in the sweep, while λ = 0.001 is indistinguishable from the RL baseline and λ = 0.01 lowers the score further.The sweep uses λ ∈ {0.001, 0.005, 0.01} on Klear-AgentForge-8B; the smallest weight is too weak to meaningfully influence training.

F.2 Multimodal Setting · G Limitations · H Broader Impacts

The multimodal ablations show that ContextRL is sensitive to contrastive-data proportion, response length, and auxiliary-loss weighting, with moderate settings performing best. The paper also limits its validation to mostly Qwen-family models below 10B parameters while targeting improved evidence-grounding reliability in agentic and visual reasoning tasks.

  • F.2 Multimodal Setting: 15% contrastive image data achieves the best overall balance, whereas increasing the proportion to 20% and especially 50% broadly degrades performance.The study evaluates 5%, 10%, 15%, 20%, and 50% proportions with Qwen2.5-VL-7B across 12 benchmarks; at 50%, nearly all benchmarks are at or below the 5% level.
  • F.2 Multimodal Setting: Maximum response length 4096 offers the best average performance, improving most benchmarks over 2048 while avoiding some 8192-token regressions.Longer responses help MathVision and MMMU-Pro, but 8192 slightly hurts ScienceQA, OlympiadBench Physics, and MME-RealWorld.
  • F.2 Multimodal Setting: Small contrastive-loss coefficients are relatively robust, but λ = 0.01 degrades performance on most benchmarks for both Qwen2.5-VL-7B and Qwen3-VL-8B.The selected defaults are λ = 0.005 for Qwen2.5-VL-7B and λ = 0.001 for Qwen3-VL-8B.
  • F.2 Multimodal Setting: The context-selection loss should augment rather than compete with GRPO, because excessive weighting can overwhelm the GRPO update signal and harm performance.The coefficient λ controls the relative weight of the context-selection loss, with λ values of 0.001, 0.005, and 0.01 compared on both models.
  • G Limitations: Due to computational constraints, experiments use base models with fewer than 10B parameters, leaving substantially larger scales and broader model-family coverage unvalidated.Most evaluated base models come from the Qwen family; the paper identifies validation at 30B+ or 70B+ and across more families as needed to strengthen generality.
  • H Broader Impacts: ContextRL explicitly rewards context grounding rather than only final answers, aiming to make multimodal and agentic LLMs more faithful to input evidence.The expected positive impact concerns improved reliability when shortcut answers are dangerous, including software-engineering agents reading code and visual reasoning systems reading fine-grained perceptual evidence.
Loading 2606.17053v1…