Context-Aware RL for Agentic and Multimodal LLMs

2026-06-15

AI Summary

Large language models often struggle to precisely locate critical evidence in long-context or complex multimodal scenarios, limiting their reasoning performance. This work proposes ContextRL, a novel approach that introduces context selection as an auxiliary proxy objective, leveraging reinforcement learning to guide the model in identifying key evidence supporting a given question-answer pair among highly similar contrastive contexts. This strategy enhances fine-grained contextual understanding. Evaluated on 1K and 7K contrastive samples derived from code agent trajectories and multimodal image scenes, ContextRL achieves average improvements of 2.2% across five long-context reasoning benchmarks and 1.8% across twelve multimodal visual question answering benchmarks, significantly outperforming standard GRPO and data augmentation baselines.

Abstract

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.

Problem

• context-awareness
• long-horizon reasoning
• multimodal reasoning
• evidence grounding
• large language models

Innovation

• Context-Aware Reinforcement Learning
• Contrastive Context Selection
• Long-Horizon Reasoning
• Multimodal LLMs
• Indirect Auxiliary Objective