Improving Multi-Step Reasoning Abilities of Large Language Models with Direct Advantage Policy Optimization

TL;DR

This work tackles the sparsity and instability challenges of RLHF in improving LLM reasoning by introducing Direct Advantage Policy Optimization (DAPO), a step-level offline RL method that uses a pretrained critic to predict per-step reasoning quality and provide dense learning signals. By decoupling the training of the actor and critic and optimizing with a quadratic loss that aligns scaled advantages with log-policy ratios, DAPO achieves stable, monotonic policy improvement under reasonable exploration assumptions. The authors provide theoretical guarantees and demonstrate substantial performance gains on mathematical and coding benchmarks across both standard and RL-finetuned models, with iterative DAPO offering further improvements. The results indicate that step-level, offline RLHF with a critic can generalize beyond the training prompts and enhance complex reasoning tasks while mitigating instability inherent in traditional actor-critic methods.

Abstract

The role of reinforcement learning (RL) in enhancing the reasoning of large language models (LLMs) is becoming increasingly significant. Despite the success of RL in many scenarios, there are still many challenges in improving the reasoning of LLMs. One challenge is the sparse reward, which makes optimization difficult for RL and necessitates a large amount of data samples. Another challenge stems from the inherent instability of RL, particularly when using Actor-Critic (AC) methods to derive optimal policies, which often leads to unstable training processes. To address these issues, we introduce Direct Advantage Policy Optimization (DAPO), an novel step-level offline RL algorithm. Unlike standard alignment that rely solely outcome rewards to optimize policies (such as DPO), DAPO employs a critic function to predict the reasoning accuracy at each step, thereby generating dense signals to refine the generation strategy. Additionally, the Actor and Critic components in DAPO are trained independently, avoiding the co-training instability observed in standard AC algorithms like PPO. We train DAPO on mathematical and code query datasets and then evaluate its performance on multiple benchmarks. Our results show that DAPO can effectively enhance the mathematical and code capabilities on both SFT models and RL models, demonstrating the effectiveness of DAPO.

Figures (2)

• Figure 1: Direct Advantage Policy Optimization (DAPO). The whole training procedure of DAPO consists of two individual stages : 1) Critic Training (left). Given the training query set $\mathcal{Q}$, for each query $q \in \mathcal{Q}$, DAPO uses the generator $\pi_{\text{gen}}$ to generate multiple rollouts from $q$ and derive a dataset of training states denoted as $\mathcal{D}_{\text{gen}}$. For each state in $s \in \mathcal{D}_{\text{gen}}$, DAPO uses the completer (often $\pi_{\text{ref}}$) to sample multiple sub-trajectory from $s$ and collects the MC value as the estimation of the true value. Then a critic network $V_\phi$ is trained to approximate the value function of the completer. 2) Policy Optimization (right). After extracting multiple next steps $\left\{ a_i \right\} _{i=1}^{n}$ of each state $s$ from the completions in step 2, DAPO then uses the trained critic to compute the advantages for all actions by $\forall i \in [n] : A_i=Q_i-\frac{1}{n}\sum_{j=1}^n{Q_j}$ and $Q_i$ is the predicted value of Concat($s,a_i$). Finally, DAPO fits the policy ratio to the advantage of each state-action pair to optimize the generation of reasoning steps.
• Figure 2: Performance improvement on MATH TEST during training process. Let $x$ be the accuracy of base model on MATH TEST and $y$ be the accuracy after DAPO training. Absolute improvement refers to $y-x$ and relative improvement refers to $\frac{y-x}{1-x}$.

Theorems & Definitions (11)

• Remark 2.1
• Lemma 2.1
• Lemma 2.2
• Theorem 2.3: Agarwal_pg_2019pg-liu
• Lemma 3.1
• Remark 3.1: Should we treat $V_\phi$ as a PRM?
• Theorem 3.2: Monotonic improvement
• Remark 3.2
• proof
• proof