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

1. Problem Statement

Visual Retrieval-Augmented Generation (Visual RAG) aims to enhance Large Vision-Language Models (LVLMs) by grounding their reasoning in external visual documents (e.g., charts, scanned reports, slides). However, existing systems face three critical limitations:

• Inaccurate Retrieval: Standard retrieval relies on generic similarity scores that fail to capture task-specific semantics, leading to the inclusion of irrelevant visual context and hallucinations.
• Ineffective Visual Utilization: Current approaches often passively encode full images, ignoring the hierarchical nature of visual understanding. This retains redundant background noise and wastes context capacity, failing to focus on information-dense regions.
• Sparse Supervision: Existing Reinforcement Learning (RL) methods typically use sparse outcome-based rewards (only at the final answer). This creates a "credit assignment problem," making it difficult to optimize intermediate decisions like retrieval, selection, and cropping.

2. Methodology: UniDoc-RL

The authors propose UniDoc-RL, a unified RL framework where an LVLM agent jointly performs retrieval, reranking, active visual perception, and reasoning within a single sequential decision-making process.

A. Hierarchical Action Space (Coarse-to-Fine)

The framework defines a "Search-Select-Perceive" action space to progressively narrow the search scope:

1. Search Action: The agent generates a text query to retrieve a broad set of candidate images from an external corpus using a standard search engine.
2. Precise Selection Action: The LVLM acts as a reranker. It analyzes the retrieved candidates and selects the most semantically relevant image(s) using <select> tags. This bridges the gap between coarse retrieval and fine-grained reasoning.
3. Visual Perception Action: Instead of processing the full image, the agent actively localizes key information. It generates bounding boxes (<bbox>) to crop and zoom into specific regions of interest (ROIs). This mimics human attention, removing redundancy and preserving high-resolution details.

B. Dense Multi-Reward Mechanism

To solve the credit assignment problem, UniDoc-RL introduces a dense reward system providing stage-specific supervision:

• Retrieval Reward ($r_{ir}$): Uses NDCG (Normalized Discounted Cumulative Gain) to evaluate the quality of the initial candidate pool.
• Selection Reward ($r_{sel}$): Uses a binary signal (1 if the selected image is relevant, 0 otherwise). To handle cases where no ground truth exists, a pseudo-supervision strategy is employed using the top-ranked candidate.
• Perception Reward ($r_{crop}$): Uses IoU (Intersection over Union) between the predicted crop and ground-truth regions to encourage precise localization.
• Pattern Reward ($r_{pat}$): A rule-based reward ensuring the output follows the correct XML-like format (e.g., <search>, <select>, <bbox>).
• Outcome Reward ($r_{ans}$): A model-based reward evaluating the correctness of the final generated answer.
• Total Reward: A weighted sum of these components ($r_{total} = \sum \lambda_i r_i$).

C. Training Paradigm

• Algorithm: The model is optimized using Group Relative Policy Optimization (GRPO), which aligns agent behavior with multiple objectives without requiring a separate value network.
• Data Synthesis: The authors curated a high-quality dataset of 12,621 samples for Supervised Fine-Tuning (SFT) and 5,537 samples for RL. They used a powerful teacher model (Qwen3-VL-235B) combined with layout analysis tools (MinerU) to synthesize reasoning trajectories with fine-grained action annotations.
• Interaction Loop: The agent iteratively interacts with the environment. Crucially, visual observations (retrieved images or cropped regions) are injected into the conversation history as "User" inputs, aligning with the model's pre-training distribution.

3. Key Contributions

1. Unified Framework: UniDoc-RL is the first framework to jointly model retrieval, reranking, active visual perception, and reasoning in a single RL process for Visual RAG.
2. Hierarchical Action Design: The introduction of "Precise Selection" and "Active Visual Perception" actions allows the model to filter noise and focus on high-value visual evidence, addressing the limitations of passive full-image encoding.
3. Dense Reward System: The multi-reward scheme provides explicit supervision for every stage of the pipeline, effectively solving the credit assignment problem in complex multi-step visual reasoning.
4. High-Quality Dataset: The release of a comprehensive dataset with fine-grained action annotations (Search, Select, Crop, Answer) provides a new benchmark for RL-based visual RAG research.

4. Experimental Results

The model was evaluated on three challenging benchmarks: SlideVQA, ViDoSeek, and MMLongBench.

• Performance: UniDoc-RL consistently outperformed state-of-the-art baselines, including Vanilla RAG, ReAct, Search-R1, and the previous RL-based method VRAG-RL.
  • On the 3B parameter model, UniDoc-RL achieved a 17.5% improvement over VRAG-RL.
  • On the 7B parameter model, it achieved a 17.7% improvement.
• Ablation Studies:
  • Removing the Selection action significantly hurt performance on reasoning-heavy tasks (SlideVQA, ViDoSeek), proving its role in context quality.
  • Removing the Visual Perception action had the largest negative impact on MMLongBench, confirming its necessity for fine-grained visual discrimination.
  • The Dense Multi-Reward system was shown to be superior to sparse outcome-only rewards, preventing error propagation in intermediate steps.
• Qualitative Analysis: RL training transformed the model from a passive reader (SFT) that often cropped entire images or irrelevant backgrounds into an active agent that precisely crops specific chart regions or table cells to extract answers.

5. Significance

UniDoc-RL represents a significant advancement in multimodal reasoning by treating visual information acquisition as an active, sequential decision-making process rather than a static retrieval task. By combining hierarchical actions with dense rewards, it demonstrates that LVLMs can learn to "look" at documents more like humans do—filtering noise, selecting relevant pages, and zooming into critical details. This approach sets a new standard for Visual RAG systems, offering a robust solution for complex document understanding tasks where precision and context management are paramount.