Offline Reinforcement Learning for LLM Multi-step Reasoning (2025)

Abstract: 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 the performance during test time.

## introduction large language models (llms) are increasingly applied to complex tasks requiring multi-step reasoning, such as mathematical problem solving @xcite @xcite , embodied agent control (wang et al., 2023; huang et al., 2022; shridhar et al., 2020; xiang * equal contribution. † corresponding author @xcite , and web navigation @xcite @xcite . enhancing llm reasoning with reinforcement learning (rl) has gained significant interest, as it offers the potential for self-improvement and learning without relying on human-labeled trajectories. however, many popular rl algorithms require costly online data collection, either by generating language onthe-fly or interacting with an environment. for instance, tuning llms with proximal policy optimization (ppo, @xcite is often prohibitively expensive for most users, which limits practical applications @xcite . in contrast, offline rl methods, such as direct preference optimization @xcite , provide a more practical approach for aligning llms with human preferences. these methods enable practitioners to tune models using pre-existing datasets, eliminating the need for live interaction or data generation. however, attempts to enhance llms' multi-step reasoning abilities with dpo are not always successful. for example, recent works find that dpo may perform close to or even worse than simpler methods like sft @xcite . additionally, dpo requires pairwise preference data. in multi-step reasoning tasks, however, data normally consists of independent trajectories with sparse rewards indicating success or failure. a common alternative is to extract correct trajectories from offline datasets and use them for supervised fine-tuning @xcite @xcite . while this approach is simple and often effective, it fails to fully exploit the offline dataset's potential-particularly the opportunity to learn from failure experience and enhance model robustness @xcite . in this paper, we introduce oreo (offline reasoning optimization), an offline rl algorithm designed to enhance llms' multi-step reasoning capabilities. building on insights from the exten-sive literature on maximum entropy rl @xcite @xcite , especially path consistency learning @xcite , oreo jointly learns a policy model and a value function by optimizing the soft bellman equation. oreo can leverage unpaired data with only sparse rewards and enables finer-grained credit assignment, which is especially critical as the correctness of reasoning trajectories often depends on a few key tokens. additionally, oreo can be extended into an iterative framework when online exploration is allowed. the trained value function is also directly available to guide the step-level beam search at inference time, further boosting performance. we demonstrate the effectiveness of our approach on both math reasoning (gsm8k, math) and embodied agent control (alfworld) tasks. it consistently outperforms baseline methods, including rejection sampling, dpo, and kto, across different model sizes. notably, we train a 1.5b model to achieve 52.5% accuracy on the math dataset using only the original training set. moreover, iterative oreo steadily improves model performance with additional training rounds, whereas baseline methods like rejection sampling exhibit signs of saturation. the value function learned by our method proves highly effective in guiding beam search for math reasoning tasks or selecting the best-of-k actions in embodied agent control. this results in up to a 17.9% relative improvement over greedy decoding on the math dataset. ## oreo: offline reasoning optimization based on the theorems presented in section 3, we present the detailed formulation of our method oreo. we further introduce two objective function variants, an iterative extension of our approach, and a test-time search strategy leveraging the value function. ## experiments in this section, we evaluate our method oreo in math reasoning and embodied agent tasks. we also demonstrate that the value function trained alongside the policy can further improve model performance at test time through step-level beam search or choosing the best-of-k action. datasets and evaluation metric. we adopt the gsm8k @xcite and math @xcite dataset for the task of math reasoning. gsm8k is a dataset of grade school math problems. it contains 7473 training problems and 1319 test problems. math consists of competition-level math problems, with a training set of size 7500 and a test set of size 5000. all problems in these datasets are labeled with step-by-step ground-truth solutions. we use the script from deepseekmath foot_0 to extract the final answer from the solution and evaluate its correctness. we adopt alfworld (shridhar et al., 2020) for the task of embodied agent control. alfworld provides interactive textworld environments for household tasks. each task is labeled with an expert trajectory. however, these data do not contain any reasoning process. @xcite annotates 3119 alfworld training trajectories with rationales for each step, allowing model training with react-style @xcite prompting. the evaluation set of alfworld contains 140 tasks in seen environments and 134 tasks in unseen environments. we evaluate the success rates of agents in completing the tasks within 40 steps. base models. for the math reasoning task, we select qwen2.5-math-1.5b @xcite and deepseekmath-7b-instruct @xcite as our base model. for the embodied agent task, we use minicpm-2b-dpo-bf16 @xcite as the base model. baseline methods. in addition to supervised finetuning, we compare our method against three other baselines: we leave more details on implementation in appendix a. ## conclusion in this paper, we present offline reasoning optimization (oreo), an offline rl algorithm for llm reasoning and embodied llm agent tasks. oreo leverages the idea of soft q-learning. it trains an explicit value function together with the llm policy by optimizing the soft bellman equation. this alleviates the need for paired preference data in dpo and enables fine-grained credit assignment among the reasoning steps. in addition, our value function can be used in test-time search to improve the model performance. we evaluate our method in gsm8k, math, and alfworld, demonstrating a consistent improvement compared to previous offline rlhf methods like dpo.