Search papers, labs, and topics across Lattice.
This paper introduces ContextRL, a context-aware reinforcement learning method designed to enhance long-horizon reasoning and multimodal performance in large language models (LLMs) by employing an indirect auxiliary objective. By rewarding models for selecting the context that best supports a given query-answer pair from two similar contexts, ContextRL encourages fine-grained grounding, which is critical for tasks that require pinpointing subtle evidence. The method demonstrates significant improvements, achieving average gains of +2.2% on long-horizon benchmarks and +1.8% on visual question answering tasks compared to standard approaches, highlighting the effectiveness of the context-selection strategy over mere data augmentation.
ContextRL reveals that fine-grained context selection can lead to substantial performance boosts in LLMs, outperforming traditional data augmentation methods.
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.
