DocPO: Advancing Document Policy Optimization via Tailored Step-Aware Rewards

2026-08-01

AI Summary

This work addresses the limited discriminative power of conventional edit-distance-based rewards in high-accuracy regions of document parsing, which hinders effective optimization. To overcome this, the authors propose DocPO, a framework incorporating a plug-in Step-Aware Annealing mechanism that progressively sharpens reward curvature during training, thereby enhancing discrimination among high-scoring samples without requiring manual supervision. The approach designs fine-grained reward signals tailored to text, tables, and mathematical expressions by integrating normalized edit distance (NED), tree-edit distance similarity (TEDS), and a hybrid Rubric+edit-distance reward. These signals are unified within a GRPO-style reinforcement learning framework. Experimental results demonstrate that DocPO significantly outperforms non-annealing baselines on the OmniDocBench and DocElemHard benchmarks, achieving consistent improvements across diverse document element types.

Abstract

Reinforcement learning (RL) for document parsing often relies on reference-based rewards rooted in edit distance (e.g., tree edit distance), yet it remains hard to optimize in the high-accuracy regime because such rewards become weakly discriminative: near-correct outputs receive very similar scores, providing limited learning signal for hard cases. We propose Step-Aware Annealing (SAA), a plug-and-play reward sharpening mechanism that progressively increases reward curvature during training, amplifying subtle quality differences among high-scoring samples while preserving stability in early learning. Built on SAA, we introduce DocPO, a document policy optimization framework with element-specific, reference-based rewards anchored by edit-distance signals: normalized string edit distance (NED) for text, tree edit distance similarity (TEDS) for tables, and a hybrid Rubric+edit reward for formulas. Experiments on OmniDocBench and DocElemHard show that SAA consistently improves GRPO-style RL across document elements over non-annealed rewards, without requiring additional human supervision for reward construction.

Problem

• document parsing
• reinforcement learning
• reward sparsity
• edit distance
• high-accuracy regime

Innovation

• Step-Aware Annealing
• Document Policy Optimization
• Reward Sharpening
• Edit Distance Rewards
• GRPO