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

Abstract: Retrieval-Augmented Generation (RAG) extends Large Vision-LLMs (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.

• The paper introduces a unified RL framework that integrates coarse retrieval, precise selection, and active visual perception to improve document reasoning.
• It employs a hierarchical action space and a dense multi-reward system to capture fine-grained semantics and optimize credit assignment.
• Empirical results show SOTA accuracy gains, outperforming baselines by up to 17.7% on visual reasoning benchmarks.

UniDoc-RL: A Hierarchical Reinforcement Learning Framework for Visual Retrieval-Augmented Generation

Motivation and Problem Statement

The paper "UniDoc-RL: Coarse-to-Fine Visual RAG with Hierarchical Actions and Dense Rewards" (2604.14967) addresses the limitations of contemporary Visual Retrieval-Augmented Generation (RAG) systems in Large Vision-LLMs (LVLMs). Existing visual RAG methodologies often treat retrieval, visual attention, and reasoning as loosely coupled, typically relying on coarse retrieval signals, static context, and outcome-based rewards. This paradigm fails to capture the fine-grained semantics and structural dependencies required for robust visual reasoning over complex documents exhibiting noise, redundancy, and ambiguous layouts.

Figure 1: Three critical factors for Visual RAG. UniDoc-RL addresses these challenges through the (a) Precise Selection action to bridge the semantic gap between coarse retrieval and reasoning, and an (b) Active Visual Perception action to focus on information-dense regions, both optimized via a (c) Dense multi-reward mechanism.

The authors argue that three factors are critical for scaling Visual RAG: accurate retrieval, fine-grained visual focus, and effective optimization strategies that resolve credit assignment throughout the multistage pipeline. The solution is instantiated as UniDoc-RL, a unified RL framework that jointly models hierarchical retrieval, reranking, active visual perception, and multi-stage reward-driven reasoning.

Architecture and Methodology

UniDoc-RL abstracts the end-to-end document reasoning process as a sequential decision-making problem, executed by an LVLM-based agent. The action space is explicitly hierarchical, instantiating a “Search-Select-Perceive” paradigm that moves from global retrieval to local perceptual grounding.

Figure 2: Overview of UniDoc-RL. (a)(b)(c) demonstrate the “Search-Select-Perceive” coarse-to-fine action space. (d) is the dense reward scheme. (e) shows the environment-agent loop and GRPO-based training.

Hierarchical Action Space

• Coarse Image Search: Initial broad candidate retrieval via external tools, framed as a search action.
• Precise Selection: LVLM-driven reranking, using learned fine-grained semantic alignment to select relevant images from the retrieved set. This bridges the semantic gap between retrieval and context for reasoning—shown to substantially increase recall.
• Active Visual Perception: Action for region cropping and zoom, focusing the agent on discriminative subregions of selected images. This mechanism deals with semantic redundancy and aligns the model with human-like attentive strategies.

Dense Multi-Reward Optimization

A significant departure from prior work is a dense, task-aware reward system:

• Image Search Reward: Leveraging NDCG to promote high-quality, ranked retrieval.
• Precise Selection Reward: Credit assignment for choosing ground-truth-relevant candidates, augmented with pseudo-positives to maintain gradient signal.
• Visual Perception Reward: IoU-based localization supervision for cropping, driving precise evidence selection.
• Pattern/Format Reward: Enforcing syntactic validity of intermediate actions in the trajectory.
• Outcome Reward: Model-based reward on the generated answer, supervising the agent toward correct end-to-end reasoning.

Training is realized via Group Relative Policy Optimization (GRPO), sidestepping separate value networks and aligning all stages of the pipeline to both sparse and dense targets.

Empirical Results

The empirical evaluation benchmarks UniDoc-RL on SlideVQA, ViDoSeek, and MMLongBench, spanning visual reasoning, multi-hop queries, and fine-grained visual discrimination. Main results demonstrate that UniDoc-RL establishes new SOTA accuracy, outperforming the strongest RL baseline VRAG-RL by 17.5% (3B model) and 17.7% (7B model) on overall accuracy.

Key Observations

• OCR-based (Text-only) RAG exhibits significantly reduced accuracy, confirming the importance of direct visual evidence and layout cues.
• Purely Visual RAG (with LVLMs) captures both textual and visual semantics, closing the gap with human-level document understanding in visually demanding cases.
• Ablations confirm that ablation of either the Precise Selection or Visual Perception actions reduces accuracy, with the Selection action yielding larger gains for multi-document/multi-hop benchmarks, and Perception being crucial for layout-intensive tasks.

Mechanistic and Behavioral Insights

A sequence of analyses illustrates critical policy improvements:

• The hierarchical reranking ("Select") action systematically improves relevant evidence recall for downstream reasoning, evidenced by improved retrieval hit rate across benchmarks.
• GRPO-based RL encourages an active cropping strategy: after RL, the agent’s invocation of “crop” actions closely matches strong teacher models and transitions from trivial/global crops to precise, localized selections.
• RL fine-tuning remediates exposure bias seen post-SFT, encouraging the model to use both the selection and visual perception tools adaptively—sometimes skipping cropping when the full image already suffices.

  Figure 3: Case studies showing iterative information synthesis, selective cropping, and adaptive tool invocation in the UniDoc-RL agent’s reasoning trajectories.

Data, Implementation, and Practical Implications

Model training is grounded in a curated, high-quality corpus combining multiple public document benchmarks with trajectories auto-generated and filtered using a teacher LVLM and expert document parser. SFT is used for initialization, and RL with dense rewards on the challenging subset further refines policy. The authors release code and data artifacts to nourish future RL-based multimodal reasoning research.

Practically, UniDoc-RL is immediately deployable as a general recipe for any LVLM-based visual agent requiring precise document question answering, evidence localization, or visual planning with complex, noisy inputs.

Theoretical and Methodological Implications

The modular, end-to-end RL architecture enables compositional reasoning over visual environments in a manner that is generalizable across data modalities, layouts, and reasoning requirements. The dense reward framework resolves credit assignment at multiple abstraction levels, a known challenge in both multimodal RL and tool-augmented agents.

The hierarchical action design is readily extensible: future research may integrate OCR-textual fusion, larger context windows, or learned environment simulators for tool abstraction. The adoption of reward models for outcome evaluation points toward tighter integration of reward learning and RL for open-ended multimodal agents.

Conclusion

UniDoc-RL substantiates a cohesive RL approach for visual RAG by explicitly modeling hierarchical acquisition of evidence and optimizing with multi-stage dense supervision. Empirically, it pushes the state-of-the-art in accuracy for visual document understanding, demonstrating clear attribution of gains to its unique selection and perception policies. The framework stands as a robust foundation for building multi-tool, adaptive multimodal agents capable of fine-grained reasoning in document-rich contexts.