PaperlayerView exact revision on arXiv

Source-linked AI summary

SWEET-RL: Training Multi-Turn LLM Agents on Collaborative Reasoning Tasks

Yifei Zhou, Song Jiang, Yuandong Tian, Jason Weston, Sergey Levine, Sainbayar Sukhbaatar, Xian Li

arXiv:2503.15478v1cs.LG

TL;DR

Multi-turn LLM-agent RL lacks effective cross-turn credit assignment while needing to preserve pretrained models’ reasoning and generalization. The paper introduces ColBench and SWEET-RL, which uses training-time information and a tailored advantage-learning objective; SWEET-RL improves ColBench success and win rates by 6% absolutely and brings Llama-3.1-8B to GPT-4o-level or better performance.

  • Problem

    Existing multi-turn RL algorithms struggle with effective credit assignment and with leveraging LLM reasoning and generalization under limited fine-tuning data.

  • Method

    The paper introduces ColBench and SWEET-RL, whose critic uses additional training-time information and whose turn-level advantage is trained with a trajectory-level Bradley-Terry objective.

  • Results

    6% absolute success and win rate gains over other SOTA algorithms enable Llama-3.1-8B to match or surpass GPT-4o and o1-mini on ColBench.

  • Takeaways & Limitations

    ColBench supports research on multi-turn RL for reasoning-intensive collaborative artifact creation, while SWEET-RL achieves improved performance in this domain.

  • Takeaways & Limitations

    The paper identifies safety concerns from agents taking over more human tasks, including malicious attacks and unexpected behaviors.

Abstract

from arXiv · show

Large language model (LLM) agents need to perform multi-turn interactions in real-world tasks. However, existing multi-turn RL algorithms for optimizing LLM agents fail to perform effective credit assignment over multiple turns while leveraging the generalization capabilities of LLMs and it remains unclear how to develop such algorithms. To study this, we first introduce a new benchmark, ColBench, where an LLM agent interacts with a human collaborator over multiple turns to solve realistic tasks in backend programming and frontend design. Building on this benchmark, we propose a novel RL algorithm, SWEET-RL (RL with Step-WisE Evaluation from Training-time information), that uses a carefully designed optimization objective to train a critic model with access to additional training-time information. The critic provides step-level rewards for improving the policy model. Our experiments demonstrate that SWEET-RL achieves a 6% absolute improvement in success and win rates on ColBench compared to other state-of-the-art multi-turn RL algorithms, enabling Llama-3.1-8B to match or exceed the performance of GPT4-o in realistic collaborative content creation.

1 Introduction

The paper introduces ColBench to evaluate multi-turn RL for realistic collaborative agent tasks and proposes SWEET-RL for improved cross-turn credit assignment. SWEET-RL uses additional training-time information and an advantage-learning objective aligned with pretrained LLMs.

  • Multi-turn objectives such as success rates are harder to optimize than next-token prediction for sequential agent decisions.
  • Existing single-turn RLHF methods lack explicit cross-turn credit assignment, while value-function methods may generalize poorly with limited fine-tuning data.
  • ColBench targets three benchmark requirements: task diversity, reasoning-intensive complexity, and minimal engineering overhead for rapid prototyping.
  • SWEET-RL gives its critic training-time information unavailable to the actor and directly models turn-level advantages instead of first learning a value function.
  • 6% absolute success and win rate gains over other SOTA algorithms let Llama-3.1-8B match or surpass GPT-4o and o1-mini on ColBench.

2 Related Work

Prior benchmarks and multi-turn RL methods leave gaps in research-friendly, reasoning-intensive evaluation and effective use of training-time information. SWEET-RL addresses these gaps through asymmetric critic observations and corresponding objective design.

  • Benchmarks for LLM Agents: Most LLM-agent benchmarks evaluate generalist models without interactive training environments and tasks for studying multi-turn RL.
  • Benchmarks for LLM Agents: Existing flexible multi-turn RL benchmarks often use narrower domains that do not require strong reasoning capabilities.
  • Multi-turn RL algorithms for LLM Agents: SWEET-RL uses additional training-time information and adjusts its optimization objective to improve multi-turn credit assignment.
  • Process reward models: Its step-wise critic resembles a process reward model, which evaluates the correctness of individual reasoning steps.

3 Collaborative Agent Benchmark (ColBench)

ColBench is designed for scalable, low-overhead evaluation of multi-turn RL on realistic collaborative artifact-creation tasks. It includes backend programming and frontend design settings with procedurally generated training data and inspected test sets.

  • Design Principles: ColBench aligns benchmark tasks with realistic artifact creation that challenges agents’ reasoning and generalization on complex requests.
  • Design Principles: A reference artifact enables LLM-simulated collaborators and objective functional evaluation, reducing setup to API access and Python packages.
  • Design Principles: ColBench contains more than 10k procedurally generated tasks and can be expanded to study scalability without overfitting.
  • Backend Programming: Backend Programming asks agents to write custom Python functions while reasoning about unspecified conditions and edge cases through clarification questions.
  • Backend Programming: Backend data generation retains functions that pass unit tests, with 10k training tasks and 1k manually inspected test tasks.
  • Frontend Design: Frontend Design asks agents to iteratively create HTML pages despite missing layout and color details supplied by a human simulator.
  • Frontend Design: Frontend Design uses 10k training tasks and 500 manually inspected test tasks, plus 6k offline training trajectories.

4 SWEET-RL: RL with Step-WisE Evaluation from Training-Time Information

SWEET-RL frames collaboration as an offline, finite-horizon POMDP and uses a two-stage procedure: directly learn a turn-wise advantage function, then use it to improve the policy. The method gives the critic additional training-time information while training the policy only on interaction history.

  • 4.2 Learning turn-wise advantage functions: SWEET-RL first trains a step-wise advantage function with additional training-time information using a Bradley-Terry objective.The advantage function is parameterized by repurposing the language-model head, with a frozen initial seed model used as the reference model.
  • 4.2 Learning turn-wise advantage functions: After critic training, the learned advantage function serves as a per-turn reward model for policy improvement using DPO.Candidate actions are sampled from the current policy and ranked by the advantage function; 16 candidates are sampled per turn, with chosen and rejected actions drawn from the top and bottom 50% quantiles.
  • 4.1 Problem Setup: Collaboration is modeled as a finite-horizon partially observable process with observable interaction history, hidden training-time information, turn-level rewards, and a cumulative-reward objective.The setting is offline because human interactions may be costly to obtain.
  • 4.2 Learning turn-wise advantage functions: Existing value-function approaches may generalize poorly with limited fine-tuning data, motivating direct learning of the turn-wise advantage function.The paper characterizes estimating advantages for each action as potentially easier for LLMs than estimating expected future returns.
  • 4.2 Learning turn-wise advantage functions: The advantage objective raises advantages for actions in chosen trajectories and lowers them for actions in rejected trajectories.The advantage is normalized by response length to stabilize training.
  • 4.2 Learning turn-wise advantage functions: The critic receives training-time information unavailable to the policy, creating an asymmetric actor-critic structure for turn-level credit assignment.The policy observes interaction history, while the advantage function can additionally condition on training-time information.

5 Experiments

The experiments evaluate SWEET-RL against LLM and multi-turn RL baselines on ColBench, then analyze credit assignment, critic design, data scaling, and off-policy transfer. SWEET-RL improves performance across collaborative programming and design tasks, with gains linked to training-time information and turn-level advantage modeling.

  • Experimental Setup: The experiments compare Llama-3.1-8B-Instruct, Llama-3.1-70B-Instruct, GPT4-O, O1-Mini, and multiple RL algorithms on ColBench.Models are evaluated in single-turn and collaborative settings, while RL comparisons use Llama-3.1-8B-Instruct.
  • Comparisons on ColBench: Multi-turn collaboration doubles success rates for all models relative to single-turn generation, but GPT4-O and O1-Mini reach only 40.4% and 30.3% success, respectively.Llama-3.1-8B-Instruct increases from 6.9% to 22.4%, while GPT4-O reaches 16.2% in the single-turn setting.
  • Comparing other algorithms with SWEET-RL: 6% and 5.4% gains over Multi-Turn DPO are reported for Backend Programming success rate and Frontend Design win rate, respectively.SWEET-RL also matches Llama-3.1-70B-Instruct with Llama-3.1-8B-Instruct on the reported comparisons.
  • Analysis: Mean-log-probability advantage parameterization outperforms natural alternatives, while removing normalization causes the actor to generate progressively shorter responses.The ablation identifies both parameterization and normalization as important design choices.
  • Analysis: SWEET-RL initially needs more data to train a reliable critic, then catches up with Multi-Turn DPO and achieves better converging performance.The reported underperformance occurs with 3k fine-tuning samples.
  • Analysis: SWEET-RL preserves a 3.8% success-rate advantage over Multi-Turn DPO with off-policy data generated by an inferior model.Using Llama-3.1-70B-Instruct, success rises from 41.8% with Multi-Turn DPO to 45.6% with SWEET-RL.

6 Conclusion

The paper introduces ColBench and SWEET-RL to advance multi-turn RL for collaborative artifact creation. Results show improved performance over state-of-the-art baselines, while the authors describe SWEET-RL as a preliminary step with substantial future research opportunities.

  • 6 Conclusion: ColBench is introduced as a low-overhead benchmark for validating multi-turn RL algorithms on reasoning-intensive collaborative artifact-creation tasks.Its design targets realistic collaborative domains and supports development of more effective multi-turn RL algorithms.
  • 6 Conclusion: SWEET-RL leverages additional training-time information and algorithmic choices to improve performance over state-of-the-art baselines on ColBench.The conclusion also reports poor generalization from off-the-shelf deep RL methods for multi-turn LLM agents.
  • 6 Conclusion: The authors characterize SWEET-RL as a preliminary step and identify better multi-turn RL algorithms as an area for future research.The stated future direction concerns effective credit assignment for multi-turn LLM agents.

Impact Statement

The impact statement connects collaborative LLMs with potential productivity gains while emphasizing safety concerns as agents assume more tasks from humans. These concerns are left for future research.

  • Impact Statement: Collaborative LLMs may improve human productivity in areas such as content creation, but agents taking over more tasks may face malicious attacks or unexpected behaviors.The authors leave these safety concerns as an important direction for future research.

A Hyperparameters

The appendix reports the hyperparameters used for SWEET-RL and baseline experiments and summarizes observed sensitivity patterns. Multi-Turn DPO and SWEET-RL are generally consistent across selected DPO hyperparameters, with a negative-log-likelihood coefficient of 0.01 often helpful.

  • A Hyperparameters: Multi-Turn DPO and SWEET-RL show consistent performance across DPO learning rate and beta settings, while a 0.01 negative-log-likelihood coefficient helps in most cases.The reported consistency concerns the DPO loss objective and the inclusion of the coefficient.
  • A Hyperparameters: Table 5 lists the hyperparameters used for SWEET-RL and baseline methods across all experiments.The table is provided for completeness and reproducibility.

B Theoretical Justifications

The section provides theoretical justification for learning advantages with trajectory-level Bradley–Terry objectives and for using critics with training-time information. It also derives supporting equalities under deterministic transitions and connects advantages to policy gradients.

  • Trajectory-level advantage learning: The trajectory-level Bradley–Terry objective can be used to learn the advantage function.The section introduces this as a theoretical justification for the method.
  • Trajectory-level advantage learning: Under deterministic transitions, the stated equality follows by telescoping the relevant terms.The proof explicitly invokes telescoping and deterministic-transition assumptions.
  • Asymmetric critic: An asymmetric critic with training-time information yields an unbiased policy-gradient estimator when averaged over samples.The critic may judge individual samples differently from its regular counterpart, while the aggregate estimator remains unbiased.
  • Policy-gradient justification: The advantage function can replace the Q-function in the policy-gradient expression.The section presents this substitution as a consequence of the preceding equality.

C Prompts

The section provides the prompts used to test different models on the Backend Programming and Frontend Programming tasks.

  • Testing prompts: The authors include the testing prompts for Backend Programming and Frontend Programming.The prompts appear in Figures 4 and 5.

D Qualitative Comparisons of Different Credit Assignment Methods

The qualitative comparison contrasts fixed LLM judges, value functions, and SWEET-RL for assigning credit to actions. SWEET-RL is described as attending to action utility and generalizing better to unseen tasks.

  • LLM-as-a-Judge: LLM-as-a-Judge can be distracted by response length and formatting instead of task-success utility.The comparison identifies this as a weakness of using a fixed judge for credit assignment.
  • Value function: A value function can assign implausibly high success probabilities to confusing early responses on unseen tasks.One example assigns the first candidate a 97% probability of final success at turn two of a ten-turn interaction.
  • SWEET-RL: SWEET-RL assigns higher advantage to a later response that identifies duplicate objects as important for task success.This contrasts with the value function’s assessment of the earlier candidate.
  • Overall comparison: SWEET can attend to the actual utility of an action for task success and generalize well, unlike the compared alternatives.The figure caption summarizes the qualitative comparison across the credit-assignment methods.

E Full Qualitative Examples

The full trajectories illustrate how SWEET-RL changes multi-turn behavior in backend programming and frontend design. Trained agents seek information, delay final answers, and use staged editing strategies, whereas zeroshot agents may conclude prematurely.

  • Trajectory overview: The qualitative trajectories are included to demonstrate ColBench task difficulty and compare different models.Backend and frontend examples cover SWEET-RL agents and zeroshot baselines.
  • Backend Programming: Zeroshot agents can ask questions but quickly conclude without enough information, producing wrong answers, including with GPT4-O.The examples indicate that this failure mode is not limited to the smaller zeroshot Llama model.
  • Backend Programming: SWEET-RL Llama-3.1-8B-Instruct learns back-and-forth information seeking before answering once the needed information is collected.The trajectories also show longer reasoning and self-corrections as emergent behaviors.
  • Frontend Design: In frontend design, SWEET-RL learns to propose a scratch solution for coarse-grained feedback before making fine-grained edits.The task requires extended-horizon reasoning over HTML code, with snippets included in each turn and trajectories reaching up to 16k tokens.
Loading 2503.15478v1…