PaperlayerView exact revision on arXiv

Source-linked AI summary

Offline Reinforcement Learning for LLM Multi-Step Reasoning

Huaijie Wang, Shibo Hao, Hanze Dong, Shenao Zhang, Yilin Bao, Ziran Yang, Yi Wu

arXiv:2412.16145v2cs.LGcs.AIcs.CL

TL;DR

Multi-step reasoning datasets often provide sparse-reward trajectories rather than the paired preferences DPO requires, making offline learning and step-level credit assignment challenging. OREO jointly learns a policy and value function through the soft Bellman Equation. It consistently outperforms offline baselines across math and embodied-control benchmarks, and its value function can guide test-time search.

  • Problem

    Multi-step reasoning commonly provides independent trajectories with sparse rewards rather than the paired preference data DPO requires, while trajectory-level optimization loses step-level credit assignment.

    Source: §3.3 · p. 4

  • Method

    OREO is an offline RL method that jointly learns an LLM policy and value function by optimizing the soft Bellman Equation.

    Source: §1 · p. 1, §6 · p. 8

  • Results

    OREO consistently outperforms offline baselines across GSM8K, MATH, and ALFWorld; a 1.5B model reaches 52.5% accuracy on MATH using the original training set.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The learned value function can guide test-time search, with up to a 17.9% relative improvement over greedy decoding on MATH.

    Source: §1 · p. 2

    Limited computation means some ablations, iterative-training, and test-time-search experiments use 1.5B models; evaluation has focused on math reasoning and embodied-agent tasks.

    Source: §7 · p. 9

Abstract

from arXiv · show

Improving the multi-step reasoning ability of large language models (LLMs) with offline reinforcement learning (RL) is essential for quickly adapting them to complex tasks. While Direct Preference Optimization (DPO) has shown promise in aligning LLMs with human preferences, it is less suitable for multi-step reasoning tasks because (1) DPO relies on paired preference data, which is not readily available for multi-step reasoning tasks, and (2) it treats all tokens uniformly, making it ineffective for credit assignment in multi-step reasoning tasks, which often come with sparse reward. In this work, we propose OREO (Offline Reasoning Optimization), an offline RL method for enhancing LLM multi-step reasoning. Building on insights from previous works of maximum entropy reinforcement learning, it jointly learns a policy model and value function by optimizing the soft Bellman Equation. We show in principle that it reduces the need to collect pairwise data and enables better credit assignment. Empirically, OREO surpasses existing offline learning methods on multi-step reasoning benchmarks, including mathematical reasoning tasks (GSM8K, MATH) and embodied agent control (ALFWorld). The approach can be extended to a multi-iteration framework when additional resources are available. Furthermore, the learned value function can be leveraged to guide the tree search for free, which can further boost performance during test time.

1 Introduction

OREO addresses the mismatch between preference-based offline learning and sparse-reward multi-step reasoning by learning a policy and value function together. It outperforms offline baselines across math and embodied-control tasks, with iterative training and value-guided search offering further gains.

1 Introduction

  • DPO requires paired preferences despite reasoning datasets commonly containing independent trajectories with sparse success-or-failure rewards.Selecting only correct trajectories for supervised fine-tuning also leaves failure experience underused.Source: §1 · p. 1
  • OREO jointly learns a policy and value function by optimizing the soft Bellman Equation, using unpaired sparse-reward data for finer-grained credit assignment.This targets reasoning settings where correctness may depend on a few key tokens.Source: §1 · pp. 1–2
  • OREO consistently outperforms rejection sampling, DPO, and KTO across GSM8K, MATH, and ALFWorld and across different model sizes.A 1.5B model reaches 52.5% accuracy on MATH using only the original training set.Source: §1 · p. 2
  • Across iterations, OREO continues improving while rejection sampling shows signs of saturation, and its value function guides test-time search.Value-guided search yields up to a 17.9% relative improvement over greedy decoding on MATH.Source: §1 · p. 2

2 Related Work

Prior LLM reinforcement-learning work spans online, rejection-sampling, preference-based, and REINFORCE-like methods, while maximum-entropy RL motivates soft Q-learning approaches. Research on LLM reasoning includes supervised and rejection-sampling methods, RL, process reward models, and test-time search; OREO connects learned values to that search direction.

2.1 Reinforcement Learning for LLM

  • LLM post-training uses PPO and also rejection-sampling, preference-based RL, and REINFORCE-like approaches.Maximum-entropy RL underlies PCL and SAC; concurrent work applies SAC to LLM reasoning, while OREO differs in derivation, KL regularization, loss variants, and empirical settings.Source: §2.1 · p. 2

2.2 LLM Reasoning

  • LLM reasoning work includes supervised fine-tuning and rejection sampling, including STaR’s generation and selection of rationales leading to correct answers.Iterative expectation-maximization approaches have also been proposed.Source: §2.2 · p. 2
  • Direct DPO applications to reasoning have not all succeeded, and collecting pairwise preference data has been described as inefficient.OREO presents its approach as addressing these DPO limitations.Source: §2.2 · p. 2
  • Process reward models provide finer-grained RL feedback, while OREO’s learned value function supports test-time scaling through tree search.The passage characterizes Monte-Carlo-rollout-trained process reward models as a special case of the learned value function.Source: §2.2 · p. 3

3 Preliminaries

The preliminaries model LLM reasoning as sequential decision-making with token actions, sparse terminal rewards, and KL regularization to a reference policy. They derive a soft Bellman basis for jointly learning policy and value, then explain how DPO’s trajectory-level preference formulation loses pairwise-data efficiency and step-level credit assignment.

3.1 MDP for LLM Reasoning

  • For text-only reasoning, each generated token is an action and the state is the prompt plus generated-token sequence; environment-interacting tasks update state after an agent response.Text-only transitions append the token to the state.Source: §3.1 · p. 3
  • The setup uses terminal-only rewards for reasoning-chain correctness or task success and KL regularization to keep the policy near a reference, typically a pretrained model followed by SFT.The reference-policy penalty is controlled by β.Source: §3.1 · p. 3

3.2 Soft Bellman Equation

  • OREO learns policy and value models by minimizing soft Bellman consistency, and satisfying that consistency for all state-action pairs guarantees optimal policy and value.The value definition incorporates KL regularization in addition to expected rewards.Source: §3.2 · pp. 3–4

3.3 Connection to DPO

  • The DPO derivation telescopes time steps, applies a Bradley-Terry preference model, and maximizes the likelihood that a preferred response beats a dispreferred one.The section introduces this derivation to explain DPO’s limitations for multi-step reasoning.Source: §3.3 · p. 4
  • DPO’s paired-data requirement is unnecessary when reasoning tasks have ground-truth rewards, and collecting such pairs inefficiently uses offline data.The Bradley-Terry preference model motivates the paired-response formulation.Source: §3.3 · p. 4
  • Relaxing soft Bellman consistency from every time step to the full trajectory loses granular credit assignment, important when a few tokens determine correctness.Source: §3.3 · p. 4

4 OREO: Offline Reasoning Optimization

OREO formulates offline LLM reasoning as soft Bellman optimization and extends the approach with loss variants, iterative training, and value-guided test-time search.

4 OREO: Offline Reasoning Optimization

  • The method section presents OREO as an offline reasoning optimization approach and introduces its objective, variants, iterative extension, and value-guided search strategy.The formulation is based on the theorems presented in Section 3.Source: §4 · p. 4

4.1 Learning Objetive

  • OREO optimizes the soft Bellman Equation for sparse-reward reasoning and trains a separate value network alongside the policy.The method enforces a telescoped version of the equation; its value network uses an MSE loss, and DPO's objective is described as a special case of the telescoped equation at t = 0.Source: §4.1 · p. 4

4.2 Loss Variants

  • Step-level OREO treats an entire reasoning step as an action, while response-level OREO enforces consistency only at the initial state.A step action's probability is the product of its token probabilities, and this variant can be grounded in the token-level MDP.Source: §4.2 · p. 5

4.3 Iterative OREO

  • OREO can be extended iteratively by collecting a new dataset with the updated policy after each training round.The new dataset contains generated responses or environment exploration and is used for further training.Source: §4.3 · p. 5

4.4 Test-Time Search with Value Function

  • A value model trained with OREO can guide beam search for math reasoning or select the highest-value sampled action for embodied tasks.The value estimates expected future reward; the paper also says its formulation removes the need for extensive heuristic designs used in prior work.Source: §4.4 · p. 5

5 Experiments

Experiments evaluate OREO on mathematical reasoning and embodied control, then test iterative training and value-guided inference. OREO outperforms baselines across the reported tasks, and its learned value function can improve test-time search.

5 Experiments

  • Evaluations cover GSM8K and MATH mathematical reasoning and ALFWorld household-task control, using step-labeled math solutions and rationale-annotated ALFWorld trajectories.Source: §5 · p. 5

5.1 Main Results

  • The response-level objective trails token-level OREO, while step-level OREO performs comparably and slightly better on GSM8K; token-level is standard for simpler implementation.The response-level variant treats actions uniformly, making sparse-reward credit assignment to individual tokens difficult.Source: §5.1 · p. 6
  • OREO outperforms all baselines in both ALFWorld settings; in unseen environments it achieves a 17.7% relative improvement over the baseline, while rejection sampling improves only marginally.Source: §5.1 · p. 6, Table 2 · p. 7
  • OREO beats all baselines across GSM8K and MATH model families; Qwen-2.5-Math 1.5B gains 5.2% relative on GSM8K and 10.5% on MATH over SFT.Source: §5.1 · p. 6, Table 1 · p. 7

5.2 Iterative OREO

  • Across three iterations, OREO steadily improves in accuracy while rejection sampling shows diminishing gains by the third round.The authors suggest OREO may benefit more by learning from failed trajectories and exploring novel failure patterns.Source: §5.2 · p. 6, Figure 2 · p. 7

5.4 Test-Time Search with Value Functions

  • 11.4% relative improvement on GSM8K and 17.9% on MATH: beam search with B = 7 outperforms greedy decoding as compute increases.Source: Figure 3 · p. 7, §5.4 · p. 8
  • In ALFWorld best-of-K selection, success rates rise rapidly and stabilize with five samples; the value function is a byproduct of OREO rather than a separately engineered component.Source: Figure 4 · p. 8, §5.4 · p. 8

5.3 Implicit vs Explicit Value Functions

  • In the challenging MATH case, Aϕ identifies the correct step where Aθ fails; for an incorrect answer, Aϕ’s value falls from 0.816 to 0.385, yielding advantage −0.431.Source: §5.3 · p. 8

6 Conclusion

OREO is an offline RL method for LLM reasoning and embodied-agent tasks that jointly trains a policy and explicit value function using the soft Bellman Equation. It reduces reliance on paired preference data, supports fine-grained credit assignment, and improves performance across GSM8K, MATH, and ALFWorld compared with prior offline RLHF methods such as DPO.

6 Conclusion

  • OREO jointly trains an LLM policy and explicit value function by optimizing the soft Bellman Equation.The method applies to LLM reasoning and embodied-agent tasks and alleviates the need for paired preference data while enabling finer-grained credit assignment among reasoning steps.Source: §6 · p. 8
  • OREO consistently improves performance on GSM8K, MATH, and ALFWorld compared with previous offline RLHF methods such as DPO.Source: §6 · p. 9
  • The learned value function can be used in test-time search to improve model performance.Source: §6 · p. 9

7 Limitations

Limited computation means some ablations, iterative OREO experiments, and test-time search experiments use 1.5B models. Evaluation has focused on mathematical reasoning and embodied-agent tasks, with coding and web browsing identified as future extensions.

7 Limitations

  • Some ablations, iterative OREO experiments, and test-time search experiments use 1.5B models because of limited computation resources.Source: §7 · p. 9
  • OREO has primarily been evaluated on mathematical reasoning and embodied-agent tasks, with coding and web browsing proposed for future evaluation.Source: §7 · p. 9

A.1 Dataset Construction

The appendix describes dataset construction, reasoning-step segmentation, and experiment hyperparameters. Training data and settings vary across tasks and model sizes.

A.1 Dataset Construction

  • For math reasoning, the appendix describes GSM8K and MATH data construction, including sampled responses, terminal rewards, and preference-pair selection for DPO.The sampling and pairing procedures differ between 1.5B and 7B models.Source: §A.1 · p. 12
  • For ALFWorld, the SFT model uses annotated data, training uses five rollouts per task, successful trajectories receive terminal reward 1, and DPO uses at most five preference pairs.Source: §A.1 · p. 12

A.2 Segmentation of Reasoning Steps

  • Step boundaries for math reasoning are indicated by line breaks in GSM8K and by line breaks and periods in MATH.GSM8K already separates reasoning steps with line breaks.Source: §A.2 · pp. 12–13

B Proof Sketch of Theorem 1

The proof sketch recasts the RLHF objective as maximum-entropy RL with a task reward plus a reference-policy log-probability term. It then applies the optimality relations and omits the discount factor used in standard RLHF settings.

B Proof Sketch of Theorem 1

  • The RLHF objective is viewed as maximum-entropy RL with reward rtask(st, at) + β log πref(at|st).Source: B Proof Sketch of Theorem 1 · p. 13
  • The proof uses prior maximum-entropy RL results stating that the optimal policy and optimal Q and value functions satisfy specified relations.The supplied passages do not state those relations explicitly.Source: B Proof Sketch of Theorem 1 · p. 13
  • In RLHF settings, the discount factor γ is normally omitted, and the proof combines the cited optimality equations using the task-plus-reference-policy reward.Source: B Proof Sketch of Theorem 1 · p. 13

C Comparison to Liu et al. (2024a)

OREO and DQO share a special-case objective equivalence but differ in their derivations, regularization, and experimental analyses. The paper focuses on math reasoning and household-simulator control and flags possible misuse risks in broader applications.

C Comparison to Liu et al. (2024a)

D Safeguard Statement

  • The paper focuses on math reasoning and household-simulator control, while noting that broader applications may pose misuse risks warranting careful safety consideration.Source: D Safeguard Statement · p. 14
Loading 2412.16145v2…