Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

RRO: LLM Agent Optimization Through Rising Reward Trajectories

Zilong Wang, Jingfeng Yang, Sreyashi Nag, Samarth Varshney, Xianfeng Tang, Haoming Jiang, Jingbo Shang, Sheikh Muhammad Sarwar

TL;DR

RRO introduces Reward Rising Optimization to efficiently scale process supervision for LLM agents by prioritizing actions that exhibit rising reward trends across trajectory steps. It combines supervised fine-tuning, reward-rising sampling, and direct preference optimization to dynamically expand next-action exploration, achieving superior performance on WebShop and InterCode-SQL with reduced exploration costs. The method is grounded in Monte Carlo process rewards and a Bellman-like progression of expected success, supported by empirical results and analysis of sampling efficiency and reward dynamics. Overall, RRO offers a practical, scalable alternative to traditional PRMs for improving multi-step reasoning in LLM-based agents.

Abstract

Large language models (LLMs) have exhibited extraordinary performance in a variety of tasks while it remains challenging for them to solve complex multi-step tasks as agents. In practice, agents sensitive to the outcome of certain key steps which makes them likely to fail the task because of a subtle mistake in the planning trajectory. Recent approaches resort to calibrating the reasoning process through reinforcement learning. They reward or penalize every reasoning step with process supervision, as known as Process Reward Models (PRMs). However, PRMs are difficult and costly to scale up with a large number of next action candidates since they require extensive computations to acquire the training data through the per-step trajectory exploration. To mitigate this issue, we focus on the relative reward trend across successive reasoning steps and propose maintaining an increasing reward in the collected trajectories for process supervision, which we term Reward Rising Optimization (RRO). Specifically, we incrementally augment the process supervision until identifying a step exhibiting positive reward differentials, i.e. rising rewards, relative to its preceding iteration. This method dynamically expands the search space for the next action candidates, efficiently capturing high-quality data. We provide mathematical groundings and empirical results on the WebShop and InterCode-SQL benchmarks, showing that our proposed RRO achieves superior performance while requiring much less exploration cost.

RRO: LLM Agent Optimization Through Rising Reward Trajectories

TL;DR

RRO introduces Reward Rising Optimization to efficiently scale process supervision for LLM agents by prioritizing actions that exhibit rising reward trends across trajectory steps. It combines supervised fine-tuning, reward-rising sampling, and direct preference optimization to dynamically expand next-action exploration, achieving superior performance on WebShop and InterCode-SQL with reduced exploration costs. The method is grounded in Monte Carlo process rewards and a Bellman-like progression of expected success, supported by empirical results and analysis of sampling efficiency and reward dynamics. Overall, RRO offers a practical, scalable alternative to traditional PRMs for improving multi-step reasoning in LLM-based agents.

Abstract

Large language models (LLMs) have exhibited extraordinary performance in a variety of tasks while it remains challenging for them to solve complex multi-step tasks as agents. In practice, agents sensitive to the outcome of certain key steps which makes them likely to fail the task because of a subtle mistake in the planning trajectory. Recent approaches resort to calibrating the reasoning process through reinforcement learning. They reward or penalize every reasoning step with process supervision, as known as Process Reward Models (PRMs). However, PRMs are difficult and costly to scale up with a large number of next action candidates since they require extensive computations to acquire the training data through the per-step trajectory exploration. To mitigate this issue, we focus on the relative reward trend across successive reasoning steps and propose maintaining an increasing reward in the collected trajectories for process supervision, which we term Reward Rising Optimization (RRO). Specifically, we incrementally augment the process supervision until identifying a step exhibiting positive reward differentials, i.e. rising rewards, relative to its preceding iteration. This method dynamically expands the search space for the next action candidates, efficiently capturing high-quality data. We provide mathematical groundings and empirical results on the WebShop and InterCode-SQL benchmarks, showing that our proposed RRO achieves superior performance while requiring much less exploration cost.

Paper Structure

This paper contains 23 sections, 9 equations, 2 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Task Formulation
  4. Outcome Supervision
  5. Process Supervision
  6. Reward Rising Optimization
  7. Supervised Fine-tuning
  8. Reward Rising Sampling
  9. Math Grounding
  10. Agent Optimization
  11. Experiments
  12. Dataset
  13. WebShop
  14. InterCode-SQL
  15. Baseline
  16. ...and 8 more sections

Figures (2)

  • Figure 1: Reward Rising Optimization: We dynamically adjust the scope of the next action exploration and prioritize action steps that exhibit a "rising reward trend" compared to their predecessors, avoiding the over-exploration or under-exploration. (a) The illustration of Direct Preference Optimization rafailov2024direct where a pair of the preference data is used to optimize the LLM agent. (b) The process reward estimation used in RRO where the average outcome reward of a set of rollouts serves as the process reward of an intermediate step. (c) The comparison of different strategies in the next action exploration stage where our RRO achieves a balance between the computation and data quality.
  • Figure 2: Sampling efficiency of RRO and IPR on WebShop and InterCode-SQL.