PaperlayerView exact revision on arXiv

Source-linked AI summary

UniDoc-RL: Coarse-to-Fine Visual RAG with Hierarchical Actions and Dense Rewards

Jun Wang, Shuo Tan, Zelong Sun, Tiancheng Gu, Yongle Zhao, Ziyong Feng, Kaicheng Yang, Zhiwu Lu

arXiv:2604.14967v2cs.CVcs.AI

TL;DR

Existing visual RAG systems often use generic retrieval signals that overlook fine-grained visual semantics needed for complex reasoning. UniDoc-RL unifies retrieval, reranking, active perception, and reasoning through hierarchical actions and dense multi-reward reinforcement learning, supported by annotated trajectories. Across three benchmarks, it consistently outperforms state-of-the-art baselines, with up to 17.7% gains over prior RL-based methods.

  • Problem

    Existing visual RAG systems typically rely on generic retrieval signals that overlook fine-grained visual semantics essential for complex reasoning.

  • Method

    UniDoc-RL uses hierarchical actions for coarse retrieval, precise selection, and active region cropping, with dense multi-reward supervision and annotated reasoning trajectories.

  • Results

    UniDoc-RL consistently outperforms state-of-the-art baselines across three benchmarks, delivering up to 17.7% gains over prior RL-based methods.

  • Takeaways & Limitations

    The framework progressively filters visual evidence and focuses perception on key regions while jointly optimizing retrieval, reranking, perception, and reasoning.

Abstract

from arXiv · show

Retrieval-Augmented Generation (RAG) extends Large Vision-Language Models (LVLMs) with external visual knowledge. However, existing visual RAG systems typically rely on generic retrieval signals that overlook the fine-grained visual semantics essential for complex reasoning. To address this limitation, we propose UniDoc-RL, a unified reinforcement learning framework in which an LVLM agent jointly performs retrieval, reranking, active visual perception, and reasoning. UniDoc-RL formulates visual information acquisition as a sequential decision-making problem with a hierarchical action space. Specifically, it progressively refines visual evidence from coarse-grained document retrieval to fine-grained image selection and active region cropping, allowing the model to suppress irrelevant content and attend to information-dense regions. For effective end-to-end training, we introduce a dense multi-reward scheme that provides task-aware supervision for each action. Based on Group Relative Policy Optimization (GRPO), UniDoc-RL aligns agent behavior with multiple objectives without relying on a separate value network. To support this training paradigm, we curate a comprehensive dataset of high-quality reasoning trajectories with fine-grained action annotations. Experiments on three benchmarks demonstrate that UniDoc-RL consistently surpasses state-of-the-art baselines, yielding up to 17.7% gains over prior RL-based methods.

1 Introduction

Visual RAG must jointly address retrieval accuracy, focused visual utilization, and stage-appropriate rewards, because existing methods often handle these factors separately. UniDoc-RL responds with hierarchical retrieval and perception actions, dense rewards, and a trajectory dataset for unified reinforcement learning.

  • Motivation: Visual RAG depends on accurate retrieval, effective visual utilization, and reasonable rewards for optimizing retrieval, selection, cropping, and reasoning.Existing approaches often address these factors in isolation or fail to model their interdependencies.
  • Limitations of Existing Methods: Generic similarity scores support coarse filtering but miss task-specific semantics and cannot adapt to query or conversational context.Passive full-image encoding also retains irrelevant background content and wastes context capacity.
  • UniDoc-RL: UniDoc-RL jointly models retrieval, selection, active visual perception, and reasoning within a unified multimodal reinforcement learning framework.Its hierarchical action space combines coarse retrieval with LVLM-driven precise selection and active crop-and-zoom operations.
  • Optimization: Dense multi-reward supervision provides stage-specific training signals throughout the decision process, addressing sparse supervision across distinct actions.The framework is presented as a unified reinforcement learning approach for visual document RAG.
  • Resources: The authors build a dataset of diverse reasoning trajectories with fine-grained action annotations for reinforcement learning in visual RAG.The dataset is identified as a main contribution alongside the unified framework and three-benchmark evaluation.

2 Related Work

Visual RAG extends retrieval-augmented generation from text toward multimodal documents, while reinforcement learning is increasingly used for LVLM reasoning and multi-step agents. Existing visual methods rely on embedding-based retrieval, and existing multimodal RL frameworks commonly use sparse outcome-based rewards.

  • Vision-based Retrieval-Augmented Generation: Traditional text-based RAG interacts with search engines, while multimodal RAG has gained attention as digital documents proliferate.Early visual RAG approaches include ColPali and VisRAG.
  • Vision-based Retrieval-Augmented Generation: Early visual RAG systems primarily use embedding-based retrieval to align textual queries with visual documents.The passage identifies ColPali and VisRAG as examples of this approach.
  • Reinforcement Learning for Multimodal Reasoning: Reinforcement learning has extended from language models to LVLMs and model-driven agents, especially for complex tasks requiring multi-step interactions.The related work passage situates this development across recent LVLM and agent frameworks.
  • Reinforcement Learning for Multimodal Reasoning: Existing reinforcement learning frameworks typically rely on sparse outcome-based rewards, which struggle to effectively guide multimodal reasoning.This limitation motivates more structured reward designs for complex multimodal tasks.

3 UniDoc-RL

UniDoc-RL treats visual information acquisition as a sequential process that progressively narrows evidence from image retrieval to semantic selection and focused perception. Its dense rewards supervise retrieval, selection, cropping, format, and final-answer quality throughout the interaction.

  • Framework: UniDoc-RL formulates retrieval, visual evidence identification, and answer generation as sequential decision-making over an image corpus.The agent iteratively generates thoughts and actions, receives observations, and stops when sufficient information has been collected.
  • Action Definition: The hierarchical action space progressively narrows visual search through Image Search, Precise Selection, and Visual Perception.Search retrieves broad candidates, selection filters semantically irrelevant images, and perception focuses on information-dense regions.
  • Action Definition: Visual Perception performs region selection, cropping, and adaptive zooming to produce query-focused observations while removing redundant content.The model specifies a target region, which the perception function converts into a high-resolution visual observation.
  • Reward Function Tailored: UniDoc-RL combines dense stage-wise feedback with pattern and model-based outcome rewards into a weighted total reward.The reward components cover retrieval, selection, cropping, valid action formatting, and final response quality.
  • Reward Function Tailored: Image Search rewards retrieval quality with NDCG over an interleaved trajectory of candidates and relevant golden images.Candidates from multiple retrieval actions are aggregated by rank before scoring against the relevant-image collection.
  • Iterative Interaction Process: The iterative interaction process executes parsed search, selection, and bounding-box actions against an external environment and appends resulting observations to the trajectory.Algorithm 1 specifies the query, policy, environment, maximum steps, action parsing, candidate retrieval, and observation construction.

4 Experiments

The experiments evaluate UniDoc-RL on three visually rich benchmarks using curated training trajectories, supervised initialization, and GRPO optimization. UniDoc-RL achieves state-of-the-art results against diverse RAG baselines, outperforming VRAG-RL by 17.5% and 17.7% on 3B and 7B models.

  • Training Data: 12,621 samples support SFT and 5,537 samples support RL after multi-stage filtering of public benchmark data.The corpus aggregates SlideVQA, Double Bench, VisR-Bench, DocBench, and DUDE.
  • Training Data: Teacher-agent trajectory synthesis uses layout analysis and Mineru-generated candidate bounding boxes to supervise retrieval, selection, and visual perception.Visual Perception actions receive intermediate layout analysis to identify potential regions of interest.
  • Training Pipeline: Supervised fine-tuning initializes reasoning and action-format adherence before GRPO fine-tuning with the multi-reward system.GRPO uses group size 5 and is run on 8 NVIDIA A100 80G GPUs.
  • Evaluation: UniDoc-RL is evaluated on ViDoSeek, SlideVQA, and MMLongBench using binary answer evaluation and overall accuracy.The benchmarks contain visually rich tasks, while the evaluator returns a binary correctness score.
  • Main Results: 17.5% and 17.7% margins over VRAG-RL are reported for UniDoc-RL with 3B and 7B models, respectively.RL-based methods outperform supervised baselines across all three datasets, while OCR-based methods lose visual cues such as spatial layout and geometric relationships.

5 Ablation Study

Ablation studies show that hierarchical actions and dense, task-aligned rewards improve UniDoc-RL, with different components contributing most strongly according to dataset characteristics.

  • Hierarchical Actions: Precise Selection substantially improves relevant-image recall and bridges the semantic gap between retrieval and reasoning.Its gains are most pronounced on reasoning-heavy SlideVQA and ViDoSeek.
  • Hierarchical Actions: Visual Perception further improves performance by actively cropping key evidence under an IoU-based reward.Its gains are larger on MMLongBench, which requires fine-grained visual discrimination.
  • Multi-Rewards: The full multi-reward framework consistently outperforms sparse outcome-based and pattern rewards across all three datasets.Dense step-wise supervision optimizes intermediate actions and guides more robust reasoning paths.
  • Multi-Rewards: Selection Reward helps most on reasoning-heavy datasets, whereas Visual Perception Reward yields the largest improvements on MMLongBench.The ablation supports aligning rewards with the specific sub-task demands of each dataset.

6 Analysis

The analysis shows that precise selection improves retrieval recall, while reinforcement learning makes visual perception more active and crops more precisely. Case studies further illustrate adaptive tool use, including when to crop and when to skip cropping.

  • 6.1 Selection Action Facilitates Better Retrieval Results.: 85.0% and 85.7% recall on SlideVQA and ViDoSeek, respectively, rose from 79.7% and 74.8% after adding precise selection.MMLongBench recall also increased from 48.9% to 52.7%.
  • 6.2 SFT Enables Tool Use, While RL Refines the Tool-Selection Behavior: 13.3% crop frequency on MMLongBench after RL training closely matched the Teacher model and exceeded SFT behavior.The result indicates learned tool invocation for detail-intensive queries.
  • 6.2 SFT Enables Tool Use, While RL Refines the Tool-Selection Behavior: RL replaced trivial full-image or background crops with precise crops that tightly bounded regions of interest.The change was guided by IoU-based perception rewards and task-specific feedback.
  • 6.3 Case Study: UniDoc-RL cropped a chart to identify a loss spike around the 150k step, but skipped cropping when a retrieved table was already legible.These cases demonstrate adaptive reasoning through RL training.

7 Conclusion

UniDoc-RL unifies retrieval, reranking, active perception, and reasoning through hierarchical actions and dense multi-reward supervision. Experiments on three benchmarks show consistent gains over state-of-the-art baselines.

  • 7 Conclusion: UniDoc-RL progressively filters visual evidence from coarse retrieval to fine-grained selection and region cropping within an LVLM agent.The framework uses GRPO for end-to-end optimization and releases reasoning trajectories with fine-grained action annotations.
  • 7 Conclusion: UniDoc-RL consistently outperforms state-of-the-art baselines on three benchmarks.

A.1 Data Collection

The data collection process aggregates diverse public benchmarks and uses a strong teacher LVLM to synthesize reasoning trajectories with coordinated retrieval, selection, and perception actions.

  • A.1 Data Collection: The training corpus spans diverse document types, languages, and reasoning requirements by aggregating samples from several public benchmarks.
  • A.1 Data Collection: The source benchmarks include SlideVQA, DoubleBench, VisRAG-Bench, DUDE, and DocBench.They cover multi-image VQA, multilingual multimodal document RAG, retrieval, multi-domain document understanding, and document reading.
  • A.1 Data Collection: A proprietary LVLM teacher synthesizes reasoning trajectories by generating thought-action pairs from the current interaction history.
  • A.1 Data Collection: For visual perception actions, Mineru detects layout elements and generates candidate bounding boxes after the teacher selects a relevant image.
  • A.1 Data Collection: The cropped view becomes the next observation, allowing the reasoning cycle to continue.

A.3 Data Filtering

The data-filtering pipeline verifies trajectory correctness, removes samples that are too easy for an intermediate SFT model, and selects challenging yet learnable RL examples.

  • A.3 Data Filtering: A multi-stage filtering pipeline is applied after trajectory synthesis to ensure data quality.
  • A.3 Data Filtering: Incorrect teacher-generated final answers are discarded through teacher-based quality verification against reference answers.
  • A.3 Data Filtering: Samples are retained as difficulty-aware examples when the teacher trajectory is correct but the intermediate SFT model answers incorrectly.This removes samples already solvable by a partially trained model.
  • A.3 Data Filtering: RL training samples are selected from five high-temperature SFT rollouts where retrieval succeeds but at least one rollout produces an incorrect answer.This targets cases where reasoning or perception, rather than retrieval, is the bottleneck.

A.4 Data Statistics

The appendix documents UniDoc-RL’s synthesized training data, adaptive visual reasoning examples, reward evaluation, and training or data-synthesis prompts. It also references the SFT and RL hyperparameter tables.

  • Data Statistics: The final dataset contains 12,621 samples for Supervised Fine-Tuning and 5,537 samples for Reinforcement Learning.These samples result from trajectory synthesis followed by multi-stage filtering, with composition detailed across source datasets in Table 4.
  • Adaptive Reasoning Examples: A representative case retrieves candidate images, selects a training-loss figure, and crops a sub-region to identify a loss spike around the 150k step.The example illustrates how active visual perception extracts fine-grained details from complex figures.
  • Reward Evaluation: The reward model evaluates a generated response against the query and reference answer, producing a binary correctness value of 0 or 1.UniDoc-RL uses Qwen2.5-72B-Instruct as the model-based reward evaluator, with its prompt shown in Figure 7.
  • Prompts and Hyperparameters: The appendix provides prompts for image search, precise selection, and visual perception data synthesis, alongside separate key-hyperparameter tables for SFT and RL.These materials are identified in Figures 9–11 and Tables 5–6.
  • Training and Testing Prompts: The training-and-testing prompt requires reasoning first and initiates image search only when necessary information is missing.Image search returns a set of images indexed starting from 0.
Loading 2604.14967v2…