UniDoc-RL: Coarse-to-Fine Visual RAG with Hierarchical Actions and Dense Rewards
| Authors | Jun Wang et al. |
| Year | 2026 |
| HF Upvotes | 10 |
| arXiv | 2604.14967 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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.
Engineering Breakdown
Plain English
This paper presents UniDoc-RL, a reinforcement learning framework that improves how large vision-language models retrieve and reason about visual information. Rather than using generic retrieval signals, the system learns to progressively refine its search from coarse-grained document retrieval down to fine-grained image selection and active region cropping. The key innovation is formulating visual information acquisition as a sequential decision-making problem with a hierarchical action space, allowing the model to suppress irrelevant content while focusing computational resources on semantically rich regions that matter for complex reasoning tasks.
Core Technical Contribution
The core novelty is replacing traditional static retrieval pipelines with a learned sequential decision-making agent that jointly optimizes four interdependent tasks: retrieval, reranking, active visual perception, and reasoning. Prior visual RAG systems treat retrieval as independent from downstream reasoning, using fixed similarity metrics that miss fine-grained visual semantics. UniDoc-RL introduces a hierarchical action space that enables progressive refinement—starting with document-level decisions and recursively zooming into image selection and region-level crops. This hierarchical formulation allows the model to learn which visual details matter for a given query, rather than blindly retrieving all potentially relevant content.
How It Works
The system operates as a reinforcement learning loop where an LVLM agent sequentially decides what to retrieve and where to focus its visual perception. First, the agent receives a query and observes available documents as state. At each step, it can take one of three hierarchical actions: retrieve a new document, select a specific image from the current document, or crop a region within an image for detailed analysis. After each action, the agent receives intermediate rewards (based on relevance and information gain) plus a final reward signal based on reasoning quality—whether the model can answer the original query correctly given the accumulated visual context. The policy is trained end-to-end using policy gradient methods, allowing credit assignment across the entire acquisition sequence so the model learns to make efficient retrieval and perception decisions.
Production Impact
Adopting UniDoc-RL would fundamentally change how visual RAG pipelines handle knowledge integration. Instead of expensive semantic similarity search across millions of documents and images, you'd deploy a learned policy that makes selective retrieval decisions, reducing both storage overhead and per-query computation by suppressing irrelevant content early. However, this introduces new production concerns: you need a reward signal during deployment (ideally from user feedback or reasoning validation), the RL policy requires careful tuning to avoid distribution shift, and inference latency becomes variable since the number of steps depends on learned decisions. The approach excels for complex visual reasoning tasks (charts, diagrams, documents) where fine-grained semantics matter, but may be overkill for simple image classification or standard QA where generic retrieval already works well. Integration would require building infrastructure to capture reward signals, monitor policy drift, and handle the variable-length action sequences.
Limitations and When Not to Use This
The paper assumes access to reward signals during training, which is realistic for human-in-the-loop setups but challenging for fully autonomous systems. The hierarchical action space design may not generalize across different document types or domains—a policy trained on technical documents might behave poorly on medical literature. Active region cropping assumes images contain localized semantically important regions; for images that are semantically homogeneous or require holistic understanding, the fine-grained refinement adds computation without benefit. The paper doesn't address scalability questions: how many documents can the agent handle before state space explosion? What's the latency penalty for RL inference versus a single deterministic retrieval pass? Additionally, the work likely requires large-scale RL training on curated datasets with ground-truth reasoning chains, making it expensive to adapt to new domains compared to simpler retrieval baselines.
Research Context
This work extends the growing intersection of reinforcement learning and retrieval-augmented generation, building on prior work in learned retrieval (e.g., DPR, ColBERT) and vision-language model augmentation (LLaVA, GPT-4V). It responds to limitations in static retrieval pipelines that treat the retrieval and reasoning stages as separate, connecting to recent work on end-to-end optimized retrieval. The hierarchical action space draws inspiration from hierarchical reinforcement learning and prior work on adaptive information acquisition in multimodal systems. This opens research directions in learned multi-stage retrieval, policy distillation for efficient inference, and zero-shot transfer of retrieval policies across domains.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::