Building Math Agents with Multi-Turn Iterative Preference Learning
Wei Xiong Chengshuai Shi Jiaming Shen Aviv Rosenberg Zhen Qin Daniele Calandriello Misha Khalman Rishabh Joshi Bilal Piot Mohammad Saleh Chi Jin Tong Zhang Tianqi Liu


Deep Brain Ultrasound Ablation Thermal Dose Modeling with in Vivo Experimental Validation
    Zhanyue Zhao, Benjamin Szewczyk, Matthew Tarasek, Charles Bales, Yang Wang, Ming Liu, Yiwei Jiang, Chitresh Bhushan, Eric Fiveland, Zahabiya Campwala, Rachel Trowbridge, Phillip M. Johansen, Zachary Olmsted, Goutam Ghoshal, Tamas Heffter, Katie Gandomi, Farid Tavakkolmoghaddam, Christopher Nycz, Erin Jeannotte, Shweta Mane, Julia Nalwalk, E. Clif Burdette, Jiang Qian, Desmond Yeo, Julie Pilitsis, and Gregory S. Fischer

Z. Zhao, B. Szewczyk, C. Bales, Y. Wang, M. Liu, Y. Jiang, K. Gandome, F. Tavakkolmoghaddam, C. Nycz, and G. S. Fischer are with Worcester Polytechnic Institute, Worcester, MA e-mail: zzhao4@wpi.edu, gfischer@wpi.edu.

B. Szewczyk, J. Qian, and J. Pilitsis are with the Department of Neurosurgery, Albany Medical Center, Albany, NY

M. Tarasek, C. Bhushan, E. Fiveland, and D. Yeo are with GE Global Research Center, Niskayuna, NY

Z. Campwala, R. Trowbridge, Z. Olmsted, S. Mane, J. Nalwalk, J. Qian, and J. Pilitsis are with the Department of Neuroscience and Experimental Therapeutics, Albany Medical Center, Albany, NY

P. M. Johansen and J. Pilitsis are with Charles E. Schmidt College of Medicine, Florida Atlantic University, Boca Raton, FL

E. Jeannotte is with Animal Resources Facility, Albany Medical Center, Albany, NY

G. Ghoshal, T. Heffter, and E. C. Burdette are with Acoustic MedSystems, Inc., Savoy, IL

This research is supported by National Institute of Health (NIH) under the National Cancer Institute (NCI) under Grant R01CA166379 and R01EB030539.

    
==========================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================








§ INTRODUCTION
 

Large language models (LLMs) have demonstrated remarkable capacities across a variety of language tasks, showcasing their broad-ranging capabilities in natural language processing. Notable models include ChatGPT <cit.>, Claude <cit.>, and Gemini <cit.>. However, despite these advances, even the most advanced closed-source LLMs still struggle with complex reasoning tasks that require multi-rounds of decision making. In particular, for the representative task of mathematical problem solving, LLMs often fail with basic arithmetic and symbolic computations <cit.>. To address this issue, recent studies recommend the integration of external tools (e.g., calculators, computational Python libraries and symbolic solvers) to augment the LLMs' mathematical problem-solving capabilities <cit.>. Specifically, by integrating natural language reasoning with the use of these external tools, these enhanced LLMs can receive external messages from tool interactions and reason based on both previously generated tokens and external messages, which significantly improves their performance in mathematical tasks <cit.>.

    These successes of tool-integrated LLMs lead to a natural research question: how can we better train LLMs to combine tool usage with intrinsic reasoning to tackle complex reasoning tasks? For the mathematical problem solving task, existing works primarily focus on synthetic data generation (by a strong teacher model) and supervised fine-tuning (SFT), as seen in ToRA <cit.>, MetaMathQA <cit.>, MAmmoTH <cit.>, and Open-MathInstruct <cit.>. These methods and synthetic datasets have yielded significant improvements in test accuracy on standard benchmarks like MATH <cit.> and GSM8K <cit.>. 

Built on strong SFT models, Reinforcement Learning from Human Feedback (RLHF) has proven to be a key technique to elicit LLMs' knowledge during the post-training stage and has become a standard practice in the LLM training pipeline <cit.>. Broadly speaking, the RLHF learning paradigm, which was originally designed for aligning large language models (LLMs) with human values and preferences <cit.>,  is distinct from SFT as it learns from relative feedback <cit.>. It has notably enhanced the capabilities of models like ChatGPT, Claude, and Gemini, enabling them to generate responses that are more helpful, harmless, and honest <cit.>. Inspired by RLHF's success in general chat applications, in this paper, we explore RLHF for improving LLMs' mathematical problem-solving abilities when equipped with external tools. In particular, since deep RL methods (e.g., the proximal policy optimization, PPO algorithm <cit.>) are often sample inefficient and unstable <cit.>, our goal is to derive direct preference learning algorithms that directly learn from the preference dataset <cit.>. 



  
Contribution. We begin by formulating the learning process as a Markov decision process (MDP), distinct from the contextual bandit approach typically used in RLHF for making general chatbots without external environment interactions <cit.>. Then, we derive the optimality condition of the optimization problem and develop multi-turn direct alignment algorithms (M-DPO and M-KTO) that incorporate external messages, where the primary modification is to mask out irrelevant tokens during training. Furthermore, we extend our approach to its online iterative variants, which recent works demonstrated to be promising <cit.>. Finally, we evaluate our approach through case studies using augmented training sets from MATH and GSM8K benchmarks, employing various base models such as Gemma <cit.>, CodeGemma <cit.>, and Mistral <cit.>. For instance, the performance of a supervised fine-tuned Gemma-1.1-it-7B model increased from 77.5% to 83.9% on GSM8K and from 46.1% to 51.2% on MATH. Similarly, a Gemma-2-it-9B model improved from 84.1% to 86.3% on GSM8K and from 51.0% to 54.5% on MATH. These empirical results indicate a significant improvement in performance over standard SFT models, demonstrating the potential of RLHF in complex reasoning task. We also provide a comprehensive recipe for the practical implementation of our online iterative multi-turn methods, and make our models, datasets, and code publicly available for further research and development.

      


 §.§ Problem Formulation
 
We denote prompt as x ∈ and assume that the interactions run for up to H rounds. At the first step, a prompt x is sampled from some distribution d_0 as the initial state s_1=x (We use the terminology “state” instead of “context” because we are concerning about an MDP instead of a contextual bandit here). Then, at each step h ∈ [H], 

    
  * Action: the agent observes the current state s_h, which is the history of the first h-1 interactions with the external environment, and takes an action a_h according to some policy π_h(·|s_h) ∈Δ(). Typically, the action is in the ReAct manner, which consist of a reasoning step f_h and an execution step e_h (e.g., writing python code) <cit.>.
    
  * Observation: in response to the agent's action, the environment then returns an observation o_h based on the history s_h and current action a_h.

Then, we transit to a new state, which is the history up to the step h+1:

    s_h+1 = (s_h, a_h, o_h) = (x, a_1, o_1, ⋯, a_h, o_h),

and a new step begins. This process repeats for H rounds in total and eventually, we collect a trajectory:

    τ = (x, a_1, o_1, ⋯, o_H-1, a_H).

See Figure <ref> for an example. The framework presented here is a Markov decision process (MDP), which offers a distinct approach from the contextual bandit model discussed in <cit.>. Formally, we define the following MDP.
 An MDP is specified by a tuple (, , H, ^*, d_0), where  is the action space, H is the episode length[In practice, the episode length can vary across the trajectories. We may additionally define that the shorter trajectories that output the final answer are in an absorbing state. We consider the fixed episode length to simplify the subsequent mathematical analysis. ], ^*={^*_h}_h=1^H are the state transition kernels, and d_0 denotes the distribution of prompt s_1=x. For each h ∈ [H], ^*_h(·|s_h,a_h) is the distribution of the next state given the state-action pair (s_h,a_h) at step h. In our setup, a trajectory τ = (x, a_1, o_1, ⋯, o_H-1, a_H) is generated by: s_1=x ∼ d_0 and for all h ∈ [H], a_h ∼π_h(·|s_h), o_h ∼^*_h(·|s_h,a_h) where s_h+1 = (s_h, a_h, o_h). When there is no ambiguity, the abbreviation s_h+1∼^*_h(·|s_h,a_h) is also adopted.










The MDP formulation of preference learning was recently studied in <cit.> but with a focus on the single-turn chat task and without explicitly considering the external messages. A unique feature of RLHF, as opposed to traditional RL studies, is the relative feedback obtained through comparisons between two trajectories that share the same initial state (prompt). We follow <cit.> to assume that the preference signal is generated by the so-called Bradley-Terry model.







 We denote τ/x = y, where the prompt is excluded from the trajectory. We assume that there exists a utility function of the trajectory u^* such that given (x, y^1, y^2), one response y^1 is preferred over another response y^2, denoted as y^1 ≻ y^2, with probability

    (y^1 ≻ y^2 | x, y^1,y^2 ) = σ(u^*(x,y^1)-u^*(x,y^2)),

where σ is the sigmoid function σ(z) = 1/(1+exp(-z)). Also, given (x, y^1, y^2) we denote the sampled preference signal as z with z=1 indicating y^1 ≻ y^2 while z=0 indicating y^2 ≻ y^1.



Under this definition, we only assume access to the trajectory-level preference, but not an action-level one. This should distinguish our approach from a straightforward extension of the single-turn RLHF <cit.>, which fixes a prompt that may include mid-trajectory steps such as (x, a_1, o_1, a_2, o_2) and look into the next single step a_3. However, we remark that the utility function itself, can be defined in a step-wise manner. To further illustrate the notion of the BT model in trajectory-level comparisons, we provide some examples of the utility function here. 


[Result Checking in Math]
Since the math reasoning datasets GSM8K <cit.> and  MATH <cit.> have the gold answer, we can check the final answer to determine the reward. In this case, u^*(x, y) = 𝕀(a_H = gold answer). 


[Outcome-supervised Reward Models (ORMs)]
Final result checking is not perfectly reliable because we can encounter false positive solutions that have the correct answer but incorrect reasoning trajectory. Instead, as shown in <cit.>, we can uniformly sample n trajectories per prompt and train an ORM to predict whether each solution is correct or not. Then, we can take the ORM prediction at the final token as the utility function. 



[Process-supervised Reward Model (PRM) and PRM without Human Annotation.]
<cit.> argues that if we can provide step-by-step supervision signal, the utility function is more effective. However, this requires more fine-grained human labels to give rating for each step of the trajectory. <cit.> studies how to automatically construct the process-labeled data for math problems with gold answers. Specifically, for s_h,a_h, we generate N trajectories with final answers [a_H^j]_j=1^N. We can define the proxy reward value: 

    r(s_h, a_h) := ∑_j=1^N 𝕀(a_H^j = gold answer) /N.

We may also use a hard version

    r(s_h, a_h) := 𝕀(There exists a  j_0: a_H^j_0 = gold answer).

Then, we can train the PRM by 

    ℒ_PRM(θ) = _τ∼[ ∑_h=1^H r(s_h,a_h) log r_θ + (1-r(s_h,a_h)) log (1-r_θ)].

In this case, we can use u^*(x, y) = min_h ∈ [H] r_θ(s_h, a_h) <cit.>, where r_θ is the constructed step-wise reward function. 


Notations.  To improve the readability of this work, we provide a notable table in Table <ref>.




 §.§ Related Work




  
LLMs for Mathematical Problem Solving. A line of works proposes to prompt LLMs to solve the complex reasoning task in a step-by-step manner, known as the Chain-of-Thought (CoT) prompting  <cit.>, which has been a standard practice in reasoning task. However, LLMs often struggle with basic arithmetic and symbolic manipulations when relying solely on internal knowledge and natural language reasoning, as measured by standard benchmarks <cit.>. To overcome these limitations, several studies have explored the use of external tools to enhance the LLMs' problem-solving abilities. This includes calculators <cit.>, symbolic solvers <cit.>, and code interpreters <cit.>. A particularly effective approach is the Program-based method (PoT), which performs CoT reasoning by writing code and using the output of the written code as the final answer <cit.>. This method significantly outperforms traditional CoT-based techniques in mathematical problem solving. However, PoT also faces challenges in planning and error handling, where natural language reasoning is more suitable <cit.>. In view of this, tool-integrated reasoning is proposed to combine the natural-language-based intrinsic reasoning with the external tools <cit.> and has achieved great progresses in recent studies <cit.>. While these efforts have primarily focused on synthetic data generation for tool-integrated reasoning, our work aims to further boost the performance of tool-integrated LLMs by RLHF.




  
RLHF and RLHF Algorithms. The predominant approach in RLHF is the deep RL method, Proximal Policy Optimization Algorithms (PPO) <cit.>, which leads to the great successes in Chat-GPT <cit.>, Gemini <cit.>, and Claude <cit.>. However, applying PPO requires extensive efforts and resources <cit.>, often beyond the scope of open-source capabilities. In view of this, alternative approaches have been developed. The rejection sampling fine-tuning was first proposed with the name RAFT (reward ranked fine-tuning) in RLHF <cit.> and was later extended to machine translation <cit.> and mathematical problem solving <cit.>. Its theoretical advantage was explored in <cit.>. Subsequently, another long line of works proposes direct preference learning algorithms, including SLiC <cit.>, DPO <cit.>, IPO <cit.>, KTO <cit.>, and GPO <cit.>. These algorithms bypass the reward modeling step and optimize carefully designed loss objectives directly on the preference dataset, hence the name direct preference learning. There are also some works focusing on more general preference structure <cit.> beyond the reward-based framework or post-processing of the model <cit.>.

The newly proposed direct preference learning algorithms have largely advanced the RLHF area, particularly the post-training of open-source models, with the Zephyr project as a notable example <cit.>. After this, a long line of work <cit.> demonstrates the effectiveness of on-policy sampling (the samples are generated by the policy to be trained) and online exploration in enhancing direct preference learning. In particular, the online iterative DPO <cit.> and its variants <cit.> have made state-of-the-art open-source models <cit.>, or even the industry models <cit.>. Despite these advancements, most algorithms are proposed and designed for single-turn interactions and chat. The scenarios beyond single-turn chat remain largely unexplored in the existing literature. One exception is the very recent work by <cit.>, which studies multi-turn chat task under general preferences. In contrast, in this paper, we aim to explore the use of RLHF in multi-turn tasks that incorporate interactions with external tools. Meanwhile, they derive a mirror-descent-based policy optimization algorithm, which is also different from ours. 



  
RLHF for Math Problem Solving. Algorithms traditionally used in general chatbot applications have been adapted to enhance the reasoning capabilities of LLMs in mathematical contexts. For instance, RAFT (Reward-rAnked Fine-Tuning) <cit.> is extensively employed for synthetic data generation, whether through on-policy (self-improving) <cit.> or off-policy (knowledge distillation) methods <cit.>. The reward signal in these scenarios is typically derived from either final result checking or Outcome-supervised Reward Models (ORMs) <cit.>. A novel approach by <cit.> introduces Process-supervised Reward Models (PRMs), which provide feedback at each step of the Chain-of-Thought, demonstrating significant improvements over ORMs when combined with rejection sampling <cit.>.

In addition to the RAFT, the GRPO algorithm proposed in <cit.> studies multi-turn math problem solving but focuses on the CoT format without external inputs and the resulting model achieves the state-of-the-art performance in its class. The GRPO is a variant of Reinforce <cit.> thus falling into the scope of deep RL methods. 

Further advancements include adapting direct preference learning algorithms to mathematical problem solving. For instance, <cit.> have applied the original DPO or KTO by taking the trajectory completion as a “meta” action.  <cit.> further adapt the online iterative DPO originally designed for chat <cit.> and achieve better performance for CoT reasoning. Inspired by the success of PRMs, recent studies have explored generating proxy step-wise labels for the intermediate steps of the reasoning trajectories. For instance, <cit.> leverage Monte Carlo Tree Search (MCTS) and use the estimated Q value to generate the proxy labels for the intermediate steps. <cit.> proposes to use AI feedback like GPT-4 <cit.> to find the first error step in the trajectory. Meanwhile, <cit.> identifies a trajectory with the correct final answer and no errors as preferable, and prompts the SFT model with a high temperature, starting from some intermediate step to collect a rejected trajectory with errors <cit.>. Finally, a very recent study by <cit.> proposes to use MCTS with a backward iteration from the final leaf node to compute the proxy unregularized value of each node. Preference pairs are then extracted from the tree by fixing the prefix and comparing the next single reasoning step. Then, they run the original DPO on these intermediate actions with the proxy labels from MCTS. To summarize, these works present different ways of preference data collection and apply the original DPO algorithm (with some additional marginal loss and regularization adapted from the literature), thereby differing from our work in both algorithmic concepts and application scope. In contrast, we study preference learning in the context of trajectory-level comparison, where we derive the optimality condition and introduce a multi-turn DPO within an online iterative framework, specifically for tool-integrated mathematical problem solving. However, we remark that while we focus on the trajectory-level comparison, the preference signal itself can be generated in a step-by-step supervision (see Section <ref> for the detailed examples). When preference signals for partial trajectories with shared prefixes are available, our method can also adapt to learn these step-level signals (see the optimality condition in (<ref>)). In particular, the algorithmic design presented in this paper can be readily combined with the MCTS-based data collection strategy outlined in recent literature, which we leave for future work.






§ ALGORITHMS DEVELOPMENT


We develop the main algorithms of this paper in this section. We proceed to handle the general MDP formulation presented in Section <ref>, which subsumes the tool-integrated mathematical reasoning problem as a special example. Therefore, the algorithms may also be applied to more general scenarios with external messages..



 §.§ Planning with a Fixed Model: Optimality Condition

Following <cit.>, we first establish the connection between any model = (, , H, , d_0, u) and its associated optimal policy. In particular, we are interested in the following KL-regularized planning problem with respect to a reference policy π_: 

    _π J(π; , π_)= _x ∼ d_0_a_h ∼π_h(·|s_h), o_h ∼_h(·|s_h,a_h)[ u(x, y) - η∑_h=1^H ( π_h(·|s_h), π_, h(·|s_h))].

In the single-turn case (i.e., H=1 and without transitions ), <cit.> show that the optimal solution with respect to a utility function u admits a closed-form solution, which is the Gibbs distribution (see Lemma <ref>):

    π_(a_1|x) ∝π_(a_1|x)exp(u(x,a_1)/η).

Moving from the single-step to multi-turn scenario, we first show that we are still concerning about the Gibbs distribution, but in a dynamic programming manner. The results are essentially from the study of entropy-regularized MDPs <cit.>. 

To illustrate the idea, we first consider the simplest case of H=2, where the model is allowed to call the tool only once. Then, our goal is to maximize the following target:

    _x ∼ d_0[ _a_1 ∼π_1(·|x)[_o_1 ∼_1(·|x, a_1)_a_2 ∼π_2(·|s_2) u(s_2, a_2) - η(π_2(·|s_2), π_, 2(·|s_2))_Inner Loop] - η(π_1(·|s_1), π_, 1(·|s_1)) ].

The idea is to take a backward iteration from h=H=2 to h=1. Specifically, when we fix s_2 and consider the inner loop, we can leverage Lemma <ref> to solve 

    π_,2(·|s_2) = _π_2_a_2 ∼π_2(·|s_2)(u(s_2,a_2) - η·(π_2(·|s_2), π_, 2(·|s_2) ) ) ∝π_, 2(·|s_2) ·exp(u(s_2,·)/η).

Then, we can define the value of the inner loop associated with π_,2 as 

    V_,2(s_2)   : = _a_2∼π_, 2 (·|s_2)[ u(s_2,a_2) - η(π_,2(·|s_2), π_, 2(·|s_2)) ]  
    
        Q_, 1(s_1,a_1)   := _o_1∼_1(·|s_1,a_1)[V_, 2(s_2)].

Then, for step h=H-1=1, we are concerning the following KL-regularized optimization problem:

    π_, 1(·|s_1) = _π_1_a_1 ∼π_1(·|x)[Q_, 1(s_1,a_1) - η(π_1(·|s_1), π_, 1(·|s_1))  ] ∝π_,1(·|s_1) ·exp(Q_, 1(s_1, ·)/η).

By construction, it can be observed that {π_, h}_h=1^2 is optimal as it maximizes the KL-regularized target. 

For general H-step MDP, we can repeat the above process for H times starting with V_, H+1 = 0 where we recursively define

    Q_, h(s_h, a_h) =    u(s_H, a_H),       if  h = H,  
       _o_h ∼_h(·|s_h, a_h) [V_, h+1(s_h+1)],      if  h ≤ H-1,

Here the optimal policy and the V-values are given by

    π_, h(a_h|s_h)    := 1/Z_h(s_h)π_, h(a_h|s_h) ·exp(Q_, h(s_h, a_h)/η)     (Gibbs distribution of Q_, h)
    
            V_, h(s_h)    := _a_h ∼π_, h(·| s_h)[Q_, h(s_h, a_h) - η·(π_, h (·|s_h), π_, h(·|s_h))]
       = ηlog_π_, h(a'_h|s_h)exp(Q_, h (s_h,a_h')/η),

where Z_h(s_h) = ∑_a_h ∈π_, h(a_h|s_h) ·exp(Q_, h(s_h, a_h)/η) is the normalization constant. The second equality in the definition of the V-value is from Lemma <ref>. Then, by definition, [π_, h]_h=1^H is the optimal policy. Essentially, we solve H Gibbs distributions in terms of the Q-values[The definitions of Q-values are different from that of <cit.> so that the optimal policy can be interpreted as the Gibbs distribution of Q-values.]. 



 §.§ Planning with a Fixed Model: Practical Algorithm

While (<ref>) can be approximately solved with standard deep RL methods, here we are interested in the implementation in a direct preference learning manner like SLiC <cit.>, DPO <cit.> or IPO <cit.>. The existing attempts <cit.> take the completion y as a “meta action” and plug it into the single-step DPO loss. In other words, they treat the external messages as the regular texts generated by the model itself. Another natural idea is to plug the probability of the trajectory into the single-step DPO loss. To be specific, for a pair (x,τ^w,τ^l), where τ^w refers to the preferred (i.e., winning) trajectory, we have

    - logσ(η[log_π(τ^l|x)/_π_(τ^l|x) - log_π(τ^w|x)/_π_(τ^w|x)] ) 
       = - logσ(η[log∏_h=1^H π_h(a^l_h|s^l_h) _h(o^l_h|s^l_h,a^l_h)/π_, h(a^l_h|s_h^l) _h(o^l_h|s^l_h,a^l_h) - log∏_h=1^H π_h(a^w_h|s^w_h) _h(o^w_h|s^w_h,a^w_h)/π_, h(a^w_h|s_h^w) _h(o^w_h|s^w_h,a^w_h)] ) 
       = -  logσ(η∑_h=1^H [logπ_h(a^l_h|s_h^l)/π_, h(a^l_h|s_h^l) - logπ_h(a^w_h|s_h^w)/π_, h(a^w_h|s_h^w)] ).

Unfortunately, the resulting algorithm does not always lead to the optimal policy as we explain next. In particular, we can solve the Q-values as 

    Q_, h(s_h, a_h)    = logπ_, h(a_h|s_h)/π_, h(a_h|s_h) + ηlog_π_, h(a'_h|s_h)exp(Q_, h (s_h,a_h')/η)
        = logπ_, h(a_h|s_h)/π_, h(a_h|s_h) + V_, h(s_h),

where two equalities uses the definition of the optimal policy π_, h and V-value V_, h in (<ref>), respectively. Furthermore, by the definition of Q-values Q_,h in (<ref>), we have

    _o_h ∼_h(·|s_h, a_h) V_, h+1(s_h+1)    = logπ_, h(a_h|s_h)/π_, h(a_h|s_h) + V_, h(s_h),     if  h ≤ H-1
    
        u(s_H, a_H)    = logπ_, H(a_H|s_H)/π_, H(a_H|s_H) + V_, H(s_H).

Summing over h ∈ [H], we have

    u(s_H, a_H)    = η∑_h=1^H logπ_, h(a_h|s_h)/π_, h(a_h|s_h) + ∑_h=1^H [V_, h(s_h) -_o_h ∼_h(·|s_h, a_h) V_, h+1(s_h+1) ]
       = η∑_h=1^Hlogπ_, h(a_h|s_h)/π_, h(a_h|s_h)_term (A) + V_, 1(s_1)_term (B) + ∑_h=1^H-1[V_, h+1(s_h+1) -_o_h ∼_h(·|s_h, a_h) V_, h+1(s_h+1) ]_term (C).

Here, term (A) is the counterpart of ηlogπ(a_1|s_1)/π_(a_1|s_1) in the single-step DPO derivation and term (B) will be cancelled if we consider the reward difference of two trajectories with the same prompt s_1 = x. Unfortunately, in practice, term (C) is typically not feasible to directly compute. Especially, some simple math with the Chebyshev's Inequality leads to that with probability at least 0.9,

    |C| ≤ 4 [∑_h=1^H-1σ_h^2]^1/2,

where σ_h^2 is the conditional variance of V_, h+1(s_h+1) -_o_h ∼_h(·|s_h, a_h) V_, h+1(s_h+1). Therefore, the bias term (C) is related to the randomness of the external environment. 

For most cases of tool-integrated LLMs for mathematical reasoning, i.e., the focus of this work, luckily the code execution result is determined by the history (the codes written by the LLMs). In other words, given the history s_h, the external observation is deterministic, which leads to term  (C)=0. Thus, with a dataset  consisting of (x, τ^w, τ^l), the following multi-turn DPO (M-DPO) loss can be adopted:

    ℒ_(θ) = -∑_(x, τ^w, τ^l) ∈logσ(η∑_h=1^H [logπ_θ, h (a^l_h|s_h^l)/π_, h(a^l_h|s_h^l) - logπ_θ, h(a^w_h|s_h^w)/π_, h(a^w_h|s_h^w)] ),
 
We emphasize again that although the loss presented in (<ref>) is identical to the one in (<ref>), a rigorous derivation procedure (rather than a direct plug-in) is provided. To the best of our knowledge, (<ref>) is new in the context of multi-turn reasoning task with external messages. In particular, it is noted that such a M-DPO loss is only valid upon deterministic transitions, i.e., term (C) = 0. 

Moreover, with (<ref>) implying that with term (C)=0, the implicit reward is given by A=η∑_h=1^Hlogπ^*_h(a_h|s_h)/π_, h(a_h|s_h), a multi-turn version of KTO <cit.>, denoted as M-KTO, can also be naturally derived:

    ℒ_(θ) = _x, y ∼[λ_y - v(x,y) ],

where 

    u_θ(x, y)    = η∑_h=1^Hlogπ_u, h(a_h|s_h)/π_, h(a_h|s_h),
    
        z_0     = _x' ∼, τ' ∼π_θ(·|x') ∑_h=1^H (π_θ(·|s_h), π_(·|s_h) ),

and

    v(x,y) = {λ_+ σ(η (u_θ(x, y) - z_0) )        if  y ∼ y_desirable | x
    λ_- σ(η (z_0 - u_θ(x, y)) )        if  y ∼ y_undesirable | x
    ..

Here λ_+ and λ_- are two hyper-parameters. We notice that <cit.> developed an online iterative version of KTO for the CoT format reasoning task. Here we extend it to build the tool-integrated reasoning agent.

The above discussions, in particular, M-DPO and M-KTO losses provided in (<ref>) and (<ref>),  are focused on deterministic observations due to the deterministic nature of tool-integrated LLMs for mathematical reasoning. In contrast, some other applications may encounter stochastic observations, e.g., multi-turn chats with the external message provided by a human or another LLM <cit.>. In these scenarios, (<ref>) is biased and cannot lead to the optimal policy since term  (C)≠ 0. Instead, one should first construct a value network based on the Bellman equations provided in  (<ref>) and (<ref>), similar to the approach in <cit.>. Subsequently, term  (C) can be estimated using Monte-Carlo methods and serve as an adaptive margin in the preference training. Particularly, the distinctions between direct preference learning algorithms and classical deep RL methods become less clear. The exploration of this more complex algorithm and its application to general multi-turn learning scenarios is left for future research.

We note that the MDP formulation above and related discussions have been previously derived by <cit.> in the context of either token-wise MDP or more general MDP with deterministic transition but their focuses are all on the single-turn chat tasks. Although the mathematical formulations appear similar, our primary focus lies on tool-integrated reasoning tasks that incorporate additional external messages {o_h}_h=1^H-1.




 §.§ Learning with Online Iterative Training



In the literature of direct preference learning, a long line of work shows that the online single-turn RLHF significantly outperforms their offline counterpart, both in the literature of direct preference learning <cit.> and DRL-based approach or rejection sampling fine-tuning <cit.>. Motivated by these successes, we propose to further incorporate online interactive learning to the multi-turn RLHF studied in this work. In the following, we illustrate the proposed ideas from mainly two aspects: two learning objectives and one unified algorithmic framework.








  
Learning objective. We consider two different learning objectives. The first one is the KL-regularized target:

    max_π_x ∼ d_0_a_h ∼π(·|s_h), o_h ∼^*_h(·|s_h,a_h)[ u^*(x, y) - η∑_h=1^H ( π(·|s_h), π_0(·|s_h))],

i.e., max_π J(π; ^*, π_0) where ^* = (, , H, ^*, d_0, u^*) is the groundtruth environment and π_0 is the initial policy (e.g., from SFT) that RLHF starts from.
This target is widely adopted in practice <cit.> and requires us to search for the optimal policy only at a fixed KL ball centered at the SFT policy π_0 <cit.>.

In contrast, the second one is the non-regularized target, i.e., directly optimizing the reward:  

    max_π_x ∼ d_0_a_h ∼π(·|s_h), o_h ∼^*_h(·|s_h,a_h)[ u^*(x, y) ] .

This target is the standard one in canonical RL studies <cit.>. One motivation for this target is that in the reasoning task, the reward function is more interpretable (e.g. final result checking) compared to the chat task. 

Additionally, we note that a stronger KL regularization in the target (<ref>) is known to be beneficial for mitigating over-fitting issue and forgetting on the out-of-domain tasks <cit.>. On the other hand, (<ref>) allows the model to move more far away, thus achieving a better in-domain performance. Thus, from one perspective, the choice between the above two targets can be viewed as a tradeoff between out-of-domain and in-domain performances. This intuition is also verified by later experiments, where optimizing the second target in (<ref>) leads to better performance on in-domain test sets. In the rest of this section, we discuss two learning objectives to fully develop the multi-turn preference learning framework. We also conduct an ablation study on these objectives in the experimental section.




  
Algorithmic framework. We present a general online iterative algorithmic framework in Algorithm <ref>. This framework is termed as Online Iterative Multi-turn Gibbs Sampling
from Human Feedback (M-GSHF) to highlight the online iterative training process and the optimal condition derived in (<ref>) that the optimal policy is a layer-wise Gibbs distribution, which generalizes the bandit formulation in <cit.>. Specifically, starting from π_0, at each iteration, we first collect a pair of trajectories by the current policy pair, where the preference signal is also revealed according to Definition <ref>. Then, we update our policy pair given the data collected so far and the next iteration begins. We now discuss some features of the framework as follows.

Policy choice for exploration-exploitation trade-off. We update our behavior policies in a non-symmetric way. The first agent, which aims to extract the historical information we have gathered so far, planning with respect to the empirically best model on the historical dataset  to get π_t^1, where the planning algorithms have been discussed in Section <ref>, e.g., optimizing the M-DPO or M-KTO loss in (<ref>) or (<ref>). However, it is widely recognized in RL studies <cit.> that simply exploiting the historical data via following the empirically best model is not sufficient to obtain a good final policy, while it is also required to explore the environment so that  new information can be collected to facilitate subsequent learning, i.e., the exploration-exploitation tradeoff. While the main agent targeting exploitation, we design the second agent, in contrast, to strategically incorporate the uncertainty of the future relative to π_t^1 given the historical information we collect so far into its policy choice. We call the policy of the second agent π_t^2 as an exploration policy because it serves to explore the underlying environment and facilitate the first agent's learning. In practice, this principle of exploration is generally interpreted as maximizing the difference between the two behavior policies or increasing the diversity of the collected data. We summarize some popular heuristic exploration policy adopted in the online iterative RLHF practice:

    
  * Mixture sampling: in the Claude project <cit.>, the authors choose to use the checkpoints from different training steps to collect data;
    
  * Inference parameters tuning: in the LLaMA project <cit.>, the authors carefully tune the sampling temperature to balance data diversity and data quality;
    
  * West-of-n sampling: <cit.> samples n responses per prompt and extract the best one and the worst one (based on some ranking criteria) to construct a preference pair.
    
We will explore the mixture sampling in the experimental section and also provide a theoretical justification in the next subsection.

Reference model choice for controlling regularization level. Despite two different learning targets are discussed in (<ref>) and (<ref>) seperately, we note that one general algorithmic framework can be adopted with the reference model choice taking as a hyper-parameter to control the regularization level and account for the two targets:

    
  * KL-regularized target in (<ref>): if we fix the reference model as the initial policy, i.e., π_t, = π_0, ∀ t∈ [T], we always search the optimal policy within the KL ball centered at π_0, and thus optimize the KL-regularized target.
    
  * Non-regularized target in (<ref>): in contrast, inspired by the mirror descent <cit.>, if we update the reference policy every iteration to be the policy learned in the last iteration, i.e., π_t,  = π^1_t-1, ∀ t∈ [T], the cumulative update can make the model to move away from the original π_0 (while a constraint is made on the per-iteration update magnitude) and we thus optimize the non-regularized target.


A graphical illustration is provided in Figure <ref> to facilitate the understanding.













 §.§ Theoretical Results

    

In this section, we show that the multi-turn RLHF problem can be solved in a statistically efficient manner under standard assumptions in learning theory literature. In particular, for generality, we target the most challenging scenario with stochastic and unknown transitions, while as aforementioned, multi-turn mathematical reasoning with external tools falls into an relatively easier regime with deterministic transitions. Here we mostly studies the KL-regularized target due to the lack of theoretical research on it. The other target of optimizing the rewards has been theoretically studied in <cit.> while the techniques of analyzing mirror-descent-style algorithm and corresponding guarantees have also be developed in <cit.>, which can be migrated to considering preference feedbacks. Also, to ease the presentation, we consider the scenario with batch size m=1, while the results can be easily generalized to large batches.

First, to measure the online learning process, we define the optimal policy as

    π^* := _π J(π):= J(π; ^*, π_0),

and introduce the standard notion of regret as

    Reg(T):=∑_t∈ [T]J(π^*) - J(π^1_t),

which represents the cumulative performance loss over T steps comparing the learned policies [π_t^1]_t=1^T against the optimal policy π^*. In addition, we consider that a bounded u^*(x, y) ∈ [0, B] for all (x, y) to maintain a reasonable utillity regime. Also,  it is assumed that we have accesses to the following policy improvement oracle, that is analogue to the one considered in <cit.>.

  For any model = (, , H, , d_0, u) and a reference function π_, we can compute the optimal policy associated with the model [π_, h]_h=1^H iteratively as in (<ref>).


The overall algorithm, i.e., the theoretical version of online iterative M-GSHF, is also summarized in Algorithm <ref>. At each round t, with = ∪_i =1^t-1_i as the aggregated dataset, it starts with performing a maximum likelihood estimation (MLE) of the reward function u^* over a set , whose elements are bounded in [0, B], as

    û_t = _û∈ L_t(û) := ∑_(x, τ^1, τ^2, z) ∈∪_i =1^t-1_i[zlog(σ(û(τ^1) - û(τ^2))) + (1-z)log(σ(û(τ^2) - û(τ^1))) ],

and also an MLE of the transition kernel ^* over a set  as

    _t = _∈ L_t() := ∑_(π, τ) ∈∪_i =1^t-1_ilog^π(τ),

where ^π(τ) denotes the probability of trajectory τ under policy π and transition kernel . With the obtained model _t = (û_t, _t), the Oracle defined in Definition <ref> is called with the reference policy π_ set as the initial policy π_0, whose output is adopted as the main policy π^1_t.

Then, we specify how to choose a theoretically sound exploration policy π_t^2. The previous work of <cit.> on single-turn RLHF has demonstrated the intuition that the exploration policy should be in charge of collecting information of the uncertain parts of the environment , which is thus often selected to maximize one uncertainty measurement. In the multi-turn RLHF setup considered in this work, the following proposition serves as the cornerstone to find a suitable uncertainty measurement to decide the exploration policy. In particular, we can observe that the optimal policy is parameterized by the optimal Q-function. If a different set of Q-function is adopted for policy parameterization, we can bound its performance as follows.

    If considering a set of Q-functions [Q̂_h]_h=1^H and a reference policy π_ with the induced policy π̂ as
    
    π̂_h(a_h|s_h) ∝π_,h(a_h|s_h) ·exp(Q̂_h(s_h, a_h)/η),

    and the corresponding set of V-functions [V̂_h]_h=1^H as
    
    V̂_h(s_h) = _a_h∼π̂_h(·|s_h)[Q̂_h(s_h, a_h)] - η(π̂_h(·|s_h), π_,h(·|s_h)),     V̂_H+1(s_H+1) = 0,

    for any comparator policy π, it holds that
    
    J(π) - J(π̂)    = _d_0,π, ^*[u^*(s_H, a_H)] - _d_0, π̂, ^*[u^*(s_H, a_H)]
       + ∑_h∈ [H]_d_0, π, ^*[V̂_h+1(s_h+1) - Q̂_h(s_h, a_h)] - ∑_h∈ [H]_d_0, π̂, ^*[V̂_h+1(s_h+1) - Q̂_h(s_h, a_h)]
       - η·∑_h∈ [H]_d_0, π, ^*[(π_h(·|s_h), π̂_h(·|s_h))],

    where the expectation _d_0, π, ^* is with respect to the prompt and response (i.e., the trajectory) generated following d_0, ^* and π.



Based on Proposition <ref>, the exploration policy π_t^2 is selected as
    
    π_t^2 = _πmax_ũ∈_t, ∈_t   _d_0, π, [ũ(s_H, a_H)] - _d_0, π^1_t, [ũ(s_H, a_H)] - (_d_0, π, [û_t(s_H, a_H)] - _d_0, π^1_t, [û_t(s_H, a_H)])_uncertainty measurement of reward estimation
       + ∑_h∈ [H]_d_0,π, [  V̂_t, h+1(s_h+1)  - [_t, hV̂_t, h+1](s_h, a_h)]_uncertainty measurement of transition estimation,

    where _t and _t are two confidence sets defined as
    
    _t    = {u∈: L_t(u) ≥ L_t(û_t) -  c_1 log(||T/δ)},
    _t    = {∈: L_t() ≥ L_t(_t) - c_1 log(||T/δ)}

with c_1 denoting an absolute constant here. Note that for the theoretical convenience, we have assumed  and  are finite here, which can be extended to the infinite case using standard discretization techniques.
It can be observed that π_t^2 is selected to maximize a combination of uncertainty from estimations of both rewards and transitions. If considering known transitions (i.e., without the need to estimate ), the uncertainty from the estimation of transitions dimimishes, which leads to a similar uncertainty measurement adopted in <cit.>.


The following theorem establishes a rigorous guarantee for the regret incurred.

    Assuming u^*∈ and ^*∈, with probability at least 1-δ, we have that
    
    (T)  ≲   κ^-1B√(d_Tlog(||T/δ)) + B^2Hξ(d_, T, c_2log(||HT/δ))
       - η·∑_h∈ [H]_d_0, π^*, ^*[(π^*_h(·|s_h), π^1_t, h(·|s_h))],

    where κ:= 1/(2+ exp(-B)+ exp(B)), c_2 is an absolute constant, d_ is the Eluder coefficient defined in Definition <ref> while d_ and ξ(·) are from the generalized Eluder-type condition defined in Definition <ref>.

We note that the Eluder coefficient  and the generalized Eluder-type condition are standard and well-adopted conditions in the theoretical studies on RL <cit.> and also RLHF <cit.>. Moreover, for a board class of RL problems (see <cit.> for more details), the Eluder coefficient d_ is small and the condition is satisfied with ξ(d_, T, c_2log(||HT/δ)) ≲√(d_Tlog(||HT/δ)), which implies that the regret of theoretical version of Algorithm <ref> is sublinear in T, further evidencing its statistical efficiency.





§ EXPERIMENTS





 §.§ Experiment Setup



  
Task, and datasets. We use the test sets of MATH <cit.> and GSM8K <cit.> to measure the model's ability to solve the mathematical problems. The MATH dataset includes 5K problems across diverse mathematical fields such as algebra, geometry, probability, number theory, and calculus. The GSM8K test set consists of 1319 grade-school math word problems, which are generally simpler than those in the MATH dataset. Examples from each dataset are as follows:

    
  * GSM8K: Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?
    
  * MATH: Find the center of the circle with equation x^2 - 6x + y^2 + 2y = 9.

To effectively solve these problems, the model needs to perform multi-turn reasoning and arithmetic operations before getting the final answer. To construct the training prompt set, we follow <cit.> to use an augmented prompt set from the 7.5K training problems of MATH and 7.47K training problems of GSM8K. In particular, we use the prompts from MetaMathQA <cit.> and MMIQC <cit.>. The new questions include rephrasing question, backward question (starting with the final answer and thinking backward to determine an unknown variable in the original question), and bootstrapping questions by in-context learning and iterative question composing <cit.>. We delete the duplicate questions and also ensure that none from the test sets of MATH and GSM8K were used. Eventually, we have 60K training prompts in total for training and randomly split them into three disjoint sets for iterative training. We also reserve a set of 1K prompts for model selection during the training.



  
Base models. We train with a range of base models, including Gemma-1.1-it-7B <cit.>, CodeGemma-1.1-it-7B <cit.>, Mistral-7B-v0.3 <cit.>, and Gemma2-it-9B. We use the pre-trained version of Mistral instead of the instruction version because the chat template of its huggingface checkpoint and that of their own code base are not consistent so we start from the pre-trained model and fine-tune it by ourselves. 




  
Data format and generation. We format the data into a multi-turn chat where the user initially ask the LLMs a question, and provide the messages returned by the Python interpreter in the subsequent user rounds of chat. In each model turn, the model reasons based the history gathered so far and can output a final answer enclosed in \boxed, or call the Python interpreter by writing a code wrapped in ```python and ```. After receiving the response of the model, we return the execution result of the code if the model calls the tool, and stop if the model outputs the final answer or reaches the maximal number of rounds H (6 in our setting). See Figure <ref> for an illustration. We generated N=30 samples per prompt for each iteration using a temperature setting of 1.0, without employing top-K or top-p sampling. We employ a mixture sampling strategy, where the up-to-date model generates only 20 trajectories, and the remainder (10 trajectories) are collected using the model from the last iteration. For the initial iteration, we employed models fine-tuned for 3 epochs and 1 epoch, respectively, to conduct mixture sampling. Intuitively, the mixture sampling helps to improve the diversity of the collected samples, and have been employed in previous RLHF practices <cit.>. For all the data generation process, we adopt the following constraints: (1) for each turn, the model can generate up to 512 tokens; (2) the maximal number of steps is H=6; (3) the maximal number of token for each trajectory is 2048. 





  
Supervised fine-tuning (SFT). We first fine-tune the model for the tool-integrated reasoning task <cit.>, using a subset of the Open-MathInstruct dataset, which was generated by the permissively licensed  model through in-context learning. The problems are from the training sets of MATH and GSM8K datasets. We restrict the number of samples for each question to be 50 and remove the nearly duplicate responses. Eventually we get 510K samples in the SFT dataset. We train the models for 4 epochs at most with a learning rate of 5e-6 for Gemma instruct models <cit.> and a learning rate of 1e-5 for Mistral-v0.3 model <cit.>. The learning rates are determined by searching {2e-6, 5e-6, 1e-5}. We use the pretrained model of Mistral because the chat template of Mistral instruct models are not consistent in different code bases (huggingface and the official one) at the time of our experiments. We use a cosine learning rate scheduler and set the warm-up steps as 100. The samples are packed into blocks with length 4096 to accelerate training and a global batch size of 64 is used. We also mask all the user messages (i.e., the prompt and the messages returned by the Python interpreter) in the training. It takes roughly 10-15 hours when training with 8xA100 80G GPUs. The checkpoint at the end of the third epoch is used for Gemma and the checkpoint of the end of the second epoch is used for Mistral as the starting point for RLHF. This is because these models outperform the last-iteration one with considerable margin and is very close to the next one.  An ablation study on the SFT epochs is also included.



  
Data Annotation. For each prompt, we first divide the responses into the winning set G^w and the losing set G^l by checking the final answer. In practice, we observe that the model can memorize the final answer and output it even though the reasoning path itself is incorrect. To mitigate this issue, we include some heuristic filtering process. First, we delete all the trajectories in the winning set where the returned messages in the second last round indicate the code is with some bugs, but the models just ignore it and predict the ground-truth answer. Then, we delete the responses in both the winning set G^w and losing set G^l if they are longer than 2048 tokens. Finally, we randomly sample a trajectory from the G^w and a trajectory from G^l to construct a pair or to add them into the training set of KTO algorithm. For each iteration, we typically get 15K-20K samples because some of the prompts may not have any correct answer. We notice that it is possible to leverage AI feedback like Gemini <cit.> or GPT4 <cit.> to further verify the correctness of the trajectory step by step or construct a PRM <cit.> to rank the trajectories, which we leave for future work.



  
Implementation of M-DPO and M-KTO. To implement the M-DPO, we simply set the labels of all the user-turn tokens to be -100 and mask the log-probability in the subsequent loss computation. We train the model for 1 epoch at most and tune the learning rate in {2e-7, 4e-7, 7e-7, 1e-6} with the first iteration of iterative training. Eventually, the learning rate of 4e-7 is used for Gemma-1.1 models and 2e-7 is used for Gemma-2 model and Mistral model. The global batch size is 32 with a warm-up step of 40. We evaluate the model every 50 training steps by the split prompt set and the best model is typically obtained between 150 steps to 600 steps, which is expected because the prompts for SFT and prompts for RLHF are overlapped. This has also been observed in previous work of RLHF for making general chatbot <cit.>. Further exploration of prompt scaling is also left for future work. The hyper-parameters are of M-KTO are mostly the same as the M-DPO. We also set the λ_+ = λ_- = 1 following the original KTO paper <cit.>. The RLHF experiments of this paper are run with 8xA100 80G GPUs, where an additional machine with 8xA100 40G GPUs is also used for data collection and model evaluation. The main experiment of this paper can be reproduced by 24 - 48 hours with this setup. We defer some other implementation details to Appendix <ref> due to space constraint.




 §.§ Main Results


We evaluate the models in the zero-shot setting and report the main results in Table <ref>. 



  
Baselines. The existing literature mainly focuses on the synthetic data generation and teach the models to use the external tool by supervised fine-tuning on the collected data. We use the results from <cit.> as baselines because we use the same SFT dataset so the results are generally comparable. For the CoT baselines, we use the Wizardmath models from <cit.>. We also include the reward ranked fine-tuning (RAFT) as a baseline <cit.>, which is also known as rejection sampling fine-tuning in the literature <cit.>. RAFT first collects N trajectories per prompt, filters the low-quality data (by reward function), and fine-tune on the selected trajectories. Another baseline is the single-turn online iterative DPO and KTO <cit.>, which ignores the problem structure (i.e., the external messages) and treats the trajectory as a whole. In implementation, it means that we do not mask the user turn and the tokens of external messages also contribute to the loss.





From the first two sections in Table <ref>, we first observe that the tool-integrated LLMs significantly outperform their CoT counterparts with only SFT, demonstrating the benefits of leveraging external tools. In the subsequent discussions, we focus on the comparison within the scope of tool-integrated LLMs. 



  
Iterative M-DPO and M-KTO considerably improve the SFT models. We observe that for all the four base models, after the iterative training with M-DPO or M-KTO, the resulting model outperforms their starting SFT checkpoint with considerable margins on both GSM8K and MATH. In particular, with M-DPO, the aligned Gemma-1.1-it-7B model attains accuracies of 83.9% and 51.2% on GSM8K and MATH, respectively, and is comparable to the open-source Open-MathInstruct-finetuned CodeLLaMA-2-70B (slightly worse on GSM8K but also slightly better on MATH). Moreover, the aligned Gemma-2-it-9B model achieves accuracies of 86.3% and 54.5% on GSM8K and MATH, surpassing all of the open-source models trained with Open-MathInstruct in the 7B to 70B range. Overall, our framework can robustly further boost the tool-integrated models' ability on the top of supervised fine-tuning. 



  
Iterative M-DPO and M-KTO surpass existing RLHF baselines. We also observe that the iterative M-DPO and M-KTO surpass other existing RLHF baselines. First, they consistently and significantly outperform the RAFT algorithm across all four base models, which is known to be a robust and competitive baseline in the literature <cit.>. This is because the RAFT algorithm only utilizes the positive signal by imitating the correct trajectories, while the DPO-based and KTO-based algorithms further leverage the negative signal from those incorrect trajectories. We note that the SFT stage in our pipeline can also be viewed as an application of RAFT, an idea that further dates back to expert iteration <cit.>. Consequently, our results should be interpreted to be that on the top of the first stage of SFT, algorithms with negative signal are more sample efficient. Moreover, while the online iterative single-turn DPO (KTO) <cit.> also gives a boost performance, it is generally worse than the multi-turn version. This suggests that learning to predict the off-policy external messages returned by the code interpreter usually has a negative impact on the reasoning ability improvement. Essentially, this corresponds to the fact that when deriving the optimality condition of the KL-regularized optimization problem, we are not allowed to optimize the external messages. Meanwhile, we present a representative example we encounter in Figure <ref>, where LLMs generate poorly constructed code resulting in anomalous and lengthy external messages. Forcing LLMs to learn to predict these messages can significantly hurt the model's reasoning abilities.












  
Iterative training and reference update lead to better performance. We use the Gemma-1.1-it-7B and M-DPO as a representative example and observe that the model benefits from online iterative training, where the test accuracy of GSM8K improves from 77.5% (SFT) to 81.5% (iter 1) to 82.5% (iter2) to 83.9% (iter3), and the test accuracy of MATH improves from 46.1% (SFT) to 49.1% (iter 1) to 49.7% (iter2) to 51.2% (iter3). This is consistent with our theoretical insight that iterative training allows the models to explore the underlying space and learn the optimal policy progressively. Moreover, we observe that if we fix the reference model as the SFT policy, the final model performance is much worse compared to the case where we update the reference model as the current model at every iteration. We suspect that this is because this version of algorithm essentially optimizes the non-regularized reward and the reward in the mathematical reasoning task is more accurate than those in the general chat task, leading to the superior in-domain performance. We defer a more detailed ablation study on the impact of KL regularization to next subsection. 






  
Preference learning improves pass@n only when n is relatively small. We plot the pass@n accuracy in terms of the number of candidate trajectories n in Figure <ref>. To evaluate the pass@n, for each question, we independently sample n trajectories, and the question is considered to be solved if there exists at least one trajectory with the correct final answer. We observe that the preference learning only improves the pass@n when n is relatively small. In particular, when n>16, all the models perform similarly on both GSM8K and MATH. In other words, the iterative M-DPO does not inject new knowledge but elicits the models' knowledge acquired in pre-training and SFT stages by boosting the quality of Top n responses. The observation is consistent with that of <cit.>, which studies the DRL-based GRPO method for the CoT mathematical reasoning task. Therefore, the success of preference learning is on top of a well-trained SFT model. We expect that the final model performance can be further improved with more high-quality SFT data.
    


 §.§ Ablation Study and Discussion

We conduct ablation studies in this subsection for a more comprehensive understanding of the proposed algorithm.



  
A moderate level of KL regularization balances the per-iteration improvement and exploration efficiency.  The effectiveness of (iterative) DPO is significantly influenced by the choice of reference model and the KL coefficient. Previous research by <cit.> on offline DPO for general chatbot applications suggests that a lower KL coefficient, specifically 0.01, yields superior performance by allowing the resulting model to move more far away from the SFT model π_0. Meanwhile, for online iterative training, <cit.> adopt a fixed reference model of π_0, and achieves continuous improvements as the training iterates. In our ablation study, we consider two different choices: (1) using the fixed reference model π_0; (2) updating the reference model to the last iteration's model at each round. Moreover, we search the KL coefficient η∈{0.01, 0.1, 0.5}. The results are summarized in Table <ref>. First, we notice that if we update the reference model at each iteration, the final model outperforms the one with a fixed reference model π_0 with a large margin. Essentially, this dynamic approach optimizes the non-regularized reward, while the one with a fixed reference model π_0 aims to maximize the KL-regularized reward. This can be viewed as a trade-off between the generation diversity and reward optimization. We suspect this performance difference is because for reasoning task, the correct reasoning paths are highly concentrated on a small subset of the generation space, and the diversity is less important in this case.

We also find that the strongest model is obtained by a moderate KL coefficient of 0.1, outperforming both 0.01 and 0.5. To understand this phenomena, we plot the test accuracy of GSM8K in Figure <ref> along the way of iterative training. As we can see, for the first iteration, the results align with <cit.>'s findings, where a smaller KL coefficient leads to a larger model improvement. However, the resulting intermediate model is further used to collect trajectories for subsequent iterative training. The models trained with very low KL coefficients tend to lose diversity rapidly, potentially reducing their capacity to collect diverse trajectories for subsequent training, leading to diminishing gains in the second and third iterations. In contrast, a higher KL coefficient of 0.5 imposes strong regularization between the resulting model and the reference model, and the model improvement is less compared to that of 0.1 for each iteration. To summarize, for online iterative training, we need to strike a balance between the per-iteration improvement and exploration efficiency to optimize the overall performance. We will see that such an intuition also extends to the choices of sampling strategy choice and other experimental tricks. 












  
The impact of sampling strategy: data diversity and coverage are crucial. Throughout our iterative training process of the Gemma-1.1-it-7B, we observed a steady increase in the percentage of correct trajectories—from 47% in the first iteration to 76% in last iteration. Moreover, since we update the reference model at each iteration, the diversity of the generated trajectories also decrease rapidly. However, the diversity of the collected data is critical for DPO/KTO training due to their contrastive nature. Prior studies in online iterative DPO for general chatbots <cit.> recommend employing model variants with different sampling temperatures or training steps to enhance trajectory diversity. Motivated by this, we explored two data collection strategies: (1) on-policy sampling, where all trajectories are sampled using the current policy, and (2) mixture sampling, where 20 trajectories are collected using the current model and 10 from the last iteration's model. We report the results in Table <ref>, where with mixture sampling, the final model performance considerably outperform the one with only on-policy sampling. To understand this phenomena, we plot the MATH test accuracy in terms of the iteration in Figure <ref>. We observe that on-policy sampling fails to improve MATH test accuracy in the third iteration, while we achieve considerable gain with the mixture sampling. This again demonstrates the importance of the diversity of the collected responses in the iterative training and also aligns with previous findings that advanced exploration strategies, which prevent diversity collapse, provide more meaningful signals for iterative preference learning <cit.>. It would be interested to explore more advanced exploration strategy like Monte Carlo tree search (MCTS) in the future study. 

In our experiments, we collected N trajectories per prompt to ensure the presence of both correct and incorrect reasoning paths for constructing the comparison pair. A larger N generally leads to a better coverage of the prompt set because for some difficult problem, we need to sample more responses to find a correct reasoning path. For instance, in iteration 1, with N=30, 92.5% of the prompts are covered, compared to 83.0% for N=12 and 60% for N=6. See Figure <ref> for an illustration of the relationship between pass@1 and N. However, increasing N also incurs higher computational costs. To understand the impact of the parameter N, we conduct an ablation study with N ∈{6, 12, 30} and summarize the results in Table <ref>. We observe a substantial performance boost when increasing N from 6 to 12, reflecting a better coverage of the complex problems that require more attempts to find a correct path. In contrast, from N=12 to N=30, we only get very minor improvement in the test accuracy, suggesting that the incremental benefits of increasing N in best-of-N sampling diminish rapidly. 




    






    


  
The best model is obtained with starting checkpoint fine-tuned with more than 1 epochs. <cit.> finds that if the SFT model is trained for more than one epoch, the subsequent DPO training will lead to performance regression with longer training in terms of instruction-following ability and benchmark for a general chatbot. In other words, there exists a trade-off between the SFT training epochs and the DPO training steps. Moreover, the best model is obtained by SFT for one epoch in their practice. We also conduct an ablation study on the impact of the SFT epoch and summarize the results in Table <ref>. Consistently across all tested scenarios, the subsequent iterative M-DPO training leads to considerable model improvement compared to the SFT model. Meanwhile, we also observe a similar trade-off between SFT and RLHF training because with more SFT epochs, the gains from the RLHF stage decrease. However, in our case, the strongest model is obtained with three epochs of SFT, followed by fine-tuning through iterative M-DPO, which is different from the offline DPO training <cit.> or the iterative DPO for general chatbot <cit.> with only one epoch of SFT.



    
    


  
NLL loss helps when the SFT model is substantially underfitting. The recent work <cit.> has introduced iterative RPO, specifically aimed at enhancing Chain of Thought (CoT) capabilities for solving mathematical problems. A key feature of this approach is the inclusion of an additional negative log-likelihood (NLL) loss for the preferred response. The main intuition for adding the NLL loss is that the original DPO algorithm <cit.> tends to reduce the likelihood of the preferred responses, and this is believed to hurt the reasoning ability <cit.>.  Motivated by their results, we explored the applicability of this idea to our setup. We conduct an ablation study by adding the NLL loss into the iterative M-DPO training and observe performance regression as reported in Table <ref>. We observe that the best model is obtained in the second iteration if we add the additional NLL loss even though we use the mixture sampling to increase the diversity of the collected data. With time-weighted exponential moving average for smoothing training record, we observe that the log probability of the chosen responses and rejected responses are (-126, -222) at the 200th step of the third iteration training when we add the NLL loss, as compared to (-166, -350) in the case without the NLL loss. This is consistent with the result of <cit.> where with the additional NLL loss, both the log probability of chosen responses and that of rejected responses increase. These evidences indicate that the NLL loss further contributes to the model distribution collapse and eventually hurt the overall performance of online iterative learning. Finally, we notice that the additional NLL loss can be viewed as an implementation of the pessimistic principle <cit.>. This also explains its inferior in-domain performance though it may be helpful to stable the training, which requires more in-depth studies.

However, one distinct feature between our setup and <cit.> is whether we first fine-tune the initialized SFT model with in-domain data. To further understand the phenomena, we fine-tune the Gemma-1.1-it-7B with only 100 steps (so that the model knows to leverage Python code to solve the problem) as the starting checkpoint of preference learning and conduct an ablation study with the NLL loss using this model. We observe when the SFT model is substantially underfitting, the addition of NLL loss actually enhances performance. This scenario mirrors the findings of <cit.>, who utilized a general LLaMA2-70B-chat model <cit.> without firstly fine-tuning on the in-domain data. Our observations align with prior research in the context of developing general chatbots <cit.>, which suggests that RLHF is less effective without preliminary SFT. 







  
On-policy sampling and small learning rate mitigate the probability drops in preferred responses. In the literature, the Direct Preference Optimization (DPO) algorithm is often reported to diminish reasoning capabilities by reducing the likelihood of preferred responses <cit.>. In our preliminary experiments, we also observe similar phenomena with a large learning rate (1e-6), where the model's reasoning ability collapses after only a few training steps, preventing convergence to good reasoning performance. In contrast, we find that using on-policy sampling within our online iterative training framework, coupled with a smaller learning rate (2e-7 or 4e-7), the DPO algorithm enhances the model's reasoning abilities. To interpret our observation, we can first write down the gradient of the DPO as follows:

    ∇_θℒ_DPO(π_θ, π_) =-η·σ( r_θ(x, y^l) - r_θ(x, y^w) ) [1/π_θ(y^w|x)∇_θπ_θ(y^w|x) - 1/π_θ(y^l|x)∇_θπ_θ(y^l|x)  ],

where r_θ(x,y) = ηlogπ_θ(x,y)/π_(x,y) is the implicit reward and we use the single-turn one for simplicity. In practice, the probability of the rejected responses typically decrease, and their gradient quickly dominates when π_θ(y^l|x) << π_θ(y^w|x) and the optimization becomes unlearning of the rejected responses. In this case, the probability of the chosen responses cannot increase. This phenomenon was also discussed in the blog <cit.>. When we adopt on-policy sampling, it leads to a relatively large probability for both rejected and chosen responses at the initial stage, ensuring that both gradients remain valid and effective. Moreover, a small learning rate prevents the model from deviating too significantly, maintaining the effectiveness of both gradients. We also notice that for the KTO algorithm, the preferred responses and the rejected responses do not appear in pairs. We suspect that the probability of the preferred response increases because the gradients of the rejected response do not dominate in every mini-batch of data. A more comprehensive understanding of the training dynamic of the direct preference learning algorithms remains largely open and we leave a more detailed study of this phenomena to future study.




§ CONCLUSION, LIMITATION, AND FUTURE RESEARCH DIRECTION


We demonstrate that preference learning, as an alternative learning paradigm to supervised fine-tuning, can further boost the performance of the tool-integrated reasoning LLMs on top of iterative best-of-n fine-tuning. We introduce an online iterative multi-turn direct preference optimization algorithm and validate its effectiveness through extensive experimentation across various base models. Our results indicate substantial improvements in the pass@1 metric over the SFT policy, as evidenced by performance gains on standard benchmarks such as GSM8K <cit.> and MATH <cit.>. Additionally, we also conduct various ablation studies to show that achieving optimal performance requires a careful balance between per-iteration improvement and exploration, facilitated by moderate levels of KL regularization and strategic exploration choices.

There are also several potential directions to further improve the model performance that we have not explored in this paper. Currently, our experiments only use final result check as the preference signal, so we cannot effectively compare trajectories that both end with correct or incorrect answers. Although our algorithm is designed for trajectory-level preference learning, the reward signal in the Bradley-Terry model could be adapted to a step-wise level. In particular, we may leverage AI feedback to verify trajectories step by step or train a process-supervised reward model <cit.> to provide learning signals. Additionally, with more fine-grained reward signals, it is also possible to adopt more advanced heuristic exploration policy like west-of-n sampling, which prove to be effective in the practice of making general chatbot <cit.> and Monte Carlo tree search (MCTS) <cit.>. Furthermore, it is also possible to leverage some well-established tricks like adaptive margin and length regularization for DPO training <cit.>. These techniques have proven to be effective for achieving a better in-domain performance for the chat task. We expect that these more fine-grained preference signals and algorithmic designs can largely improve the models' performance.

Finally, while the direct preference learning algorithms show promising gains for the mathematical reasoning tasks with code interpreter, it is not directly applicable to the general agent learning with more complex and stochastic external environments or against dynamic opponents. In particular, it requires to construct a value network for involving an adaptive margin in the optimization target and take the randomness of the external environment into consideration. We leave the study of this more involved algorithm to the future work. Moving beyond the framework presented this paper, it is also possible to explore more general preference structures beyond the BT model <cit.>. We hope that the insights from this paper will inspire further research in this direction, extending the utility of preference learning beyond the general structured chat tasks.










abbrvnat
*








§ NOTATION TABLE








§ IMPLEMENTATION DETAIL





  
Tools in Math Problem Solving.
Following <cit.>, the LLM agent is allowed to call the python interpreter when it decodes a python code starting with ```python and ending with ```. For each step h, to generate the observation o_h, we leverage the python package , and run all the codes in the history one by one and treat each code snippet as a Jupyter cell. We only return the standard output or the error message from the last snippet. When there exists some bug in the code, we only return the error message which is typically less than 20 tokens as in <cit.>. We notice that some works (e.g. <cit.>) also returns the first and the last 50 tokens of the trackback information. 
  



  
Data Generation. All the models are evaluated in the zero-shot setting. For all the data generation process, we adopt the following constraints: (1) for each turn, the model can generate up to 512 tokens; (2) the maximal number of steps is H=6; (3) the maximal number of generated token for each trajectory is 2048. When collecting new data for online iterative M-DPO, we set temperature to be 1.0 and decode without top-K or top-p sampling. For evaluation, greedy decoding is employed so that the results are generally comparable with previous works <cit.>. For evaluating the models with pass@n rate, we follow <cit.> to adopt a temperature of 0.7.



  
Python Experiment Environment. We find that the evaluation can be influenced by the python environment, the precision (especially for the Gemma-1.1 models), and even the virtual machine we use. This does not affect the overall trend and conclusion because the magnitude of oscillation is relatively small compared to the overall improvement. For completeness, however, we specify some of the key package versions here. We use transformers 4.42.4, torch 2.3.0, sympy 1.2, antlr4-python3-runtime 4.11.0, IPython 8.26.0 for all models. We evaluate the models using torch.float and use vllm 0.5.0.post1 for most the experiments except for Gemma-2 where vllm 0.5.1 is required. The inconsistency of vllm version is because Gemma-2 model was not released when we performed the main experiments of this project. We fix the python environment and machine for our evaluation throughout the experiment. For SFT, we use the open-source axolotl project with version 0.4.1 and for online iterative preference learning and RAFT, we use the code base from RLHF Workflow <cit.>. 



  
RAFT implementation. The data generation step is similar to the online iterative M-DPO training, except that we only keep the trajectories with correct final answer. For each prompt, we sample at most k trajectories where we search k ∈{1, 3, 8} and use k=1 eventually because we do not see improvement by leveraging more data. We run the algorithm for three iterations in total. The training parameters are similar to the SFT stage, but we use a smaller batch size of 32 so that there are enough optimization steps. For Gemma models, we use a learning rate of 5e-6. For each training stage, we train the models for two epochs in total according to our parameter search. For Mistral model, we find that a smaller learning rate of 1e-6 and training for 1 epoch give us much better performance. 




  
Prompt template. We do not tune the prompt though we do observe that the prompt engineering can further improve the performance. For all the experiments, we simply adopt the chat template of the models as in Figure <ref>.







§ OMITTED THEORETICAL PROOFS



 §.§ Proof of Proposition <ref>


For one policy π, starting with V^π_, H+1 = 0, we recursively define its V-value and Q-value functions on one model = (, , H, , d_0, u) and the reference policy π_ as

    Q^π_, h(s_h, a_h)    :=    u(s_H, a_H),       if  h = H,  
       _o_h ∼_h(·|s_h, a_h) [V^π_, h+1(s_h+1)],      if  h ≤ H-1, 
    
    
        V^π_, h(s_h)    := _a_h ∼π_h(·| s_h)[Q^π_, h(s_h, a_h) - η·(π_h (·|s_h), π_, h(·|s_h))].

It is noted that with the optimal policy π_, Q_, h = Q^π__, h and V_, h = V^π__, h. In the following discussions, we exclusively focus on the model ^* = (, , H, ^*, d_0, u^*) with abbreviations Q^π_h = Q^π_^*, h and  V^π_h = V^π_^*, h.

For any comparator policy π, it holds that

    J(π) - J(π̂)  = _d_0[V_1^π(s_1) - V̂_1(s_1)] - _d_0[V_1^π̂(s_1) - V̂_1(s_1)],

For any h∈ [H], we can obtain that

    _d_0, π_1:h-1, ^*_1:h-1[V_h^π(s_h) - V̂_h(s_h)] - _d_0, π̂_1:h-1, ^*_1:h-1[V_h^π̂(s_h) - V̂_h(s_h)]
       (a)=_d_0, π_1:h-1, ^*_1:h-1[_π_h[Q_h^π(s_h, a_h)] - η(π_h(·|s_h),π_,h(·|s_h))]  
        - _d_0, π_1:h-1, ^*_1:h-1[ _π̂_h[Q̂_h(s_h, a_h)] - η(π̂_h(·|s_h),π_,h(·|s_h))]
       - _d_0, π̂_1:h-1, ^*_1:h-1[_π̂_h[Q_h^π̂(s_h, a_h)] - η(π̂_h(·|s_h), π_,h(·|s_h))]
       + _d_0, π̂_1:h-1, ^*_1:h-1[_π̂_h[Q̂_h(s_h, a_h)]- η(π̂_h(·|s_h), π_,h(·|s_h))]
        = _d_0, π_1:h, ^*_1:h-1[Q_h^π(s_h, a_h) - Q̂_h(s_h, a_h)] - _d_0, π̂_1:h, ^*_1:h-1[Q_h^π̂(s_h, a_h) - Q̂_h(s_h, a_h)]
        + _d_0, π_1:h-1, ^*_1:h-1[_π_h[Q̂_h(s_h, a_h)] - _π̂_h[Q̂_h(s_h, a_h)]]_term (I)
        -  η·_d_0, π_1:h-1, ^*_1:h-1[(π_h(·|s_h), π_,h(·|s_h))] +  η·_d_0, π_1:h-1, ^*_1:h-1[(π̂_h(·|s_h), π_,h(·|s_h))]
       (b)=_d_0, π_1:h, ^*_1:h-1[Q_h^π(s_h, a_h) - Q̂_h(s_h, a_h)] - _d_0, π̂_1:h, ^*_1:h-1[Q_h^π̂(s_h, a_h) - Q̂_h(s_h, a_h)]
        -  η·_d_0, π_1:h-1, ^*_1:h-1[(π_h(·|s_h), π̂_h(·|s_h))].

In the above derivation, equation (a) is from the definitions of Q^π and V^π, and the relationship between Q̂ and V̂. The equation (b) is because

    (term I)   := _π_h[Q̂_h(s_h, a_h)] - _π̂_h[Q̂_h(s_h, a_h)] 
        = η·_π_h[logπ̂_h(a_h|s_h)/π_,h(a_h|s_h)] -  η·_π̂_h[logπ̂_h(a_h|s_h)/π_,h(a_h|s_h)]
        =  η·(π_h(·|s_h),π_,h(·|s_h)) -  η·(π_h(·|s_h), π̂_h(·|s_h)) -  η·(π̂_h(·|s_h), π_,h(·|s_h)).

where the second equation is from the relationship that

    Q̂_h(s_h, a_h) = η·logπ̂_h(a_h|s_h)/π_,h(a_h|s_h) - η·logẐ_h(s_h).


Furthermore, if h = H, we can obtain that

    _d_0, π_1:H-1, ^*_1:H-1[V_H^π(s_H) - V̂_H(s_H)] - _d_0, π̂_1:H-1, ^*_1:H-1[V_H^π̂(s_H) - V̂_H(s_H)]
        = _d_0, π_1:H, ^*_1:H-1[u^*(s_H, a_H) - Q̂_H(s_H, a_H)]  - _d_0, π̂_1:H, ^*_1:H-1[u^*(s_H, a_H) - Q̂_H(s_H, a_H)]
        -  η·_d_0, π_1:H-1, ^*_1:H-1[(π_H(·|s_H), π̂_H(·|s_H))] 
        = _d_0, π_1:H, ^*_1:H-1[u^*(s_H, a_H)] - _d_0, π̂_1:H, ^*_1:H-1[u^*(s_H, a_H)]
        + _d_0, π_1:H, ^*_1:H[V̂_H+1(s_H+1) - Q̂_H(s_H, a_H)]  - _d_0, π̂_1:H, ^*_1:H[V̂_H+1(s_H+1) - Q̂_H(s_H, a_H)]
        -  η·_d_0, π_1:H-1, ^*_1:H-1[(π_H(·|s_H)||π̂_H(·|s_H))],

where the second equality leverages that V̂_H+1(s_H+1) = 0;
otherwise, for all h ≤ H-1, it holds that

    _d_0, π_1:h-1, ^*_1:h-1[V_h^π(s_h) - V̂_h(s_h)] - _d_0, π̂_1:h-1, ^*_1:h-1[V_h^π̂(s_h) - V̂_h(s_h)]
        = _d_0, π_1:h, ^*_1:h-1[Q^π_h(s_h,a_h) - Q̂_h(s_h, a_h)] - _d_0, π̂_1:h, ^*_1:h-1[Q^π̂_h(s_h,a_h) - Q̂_h(s_h, a_h)]
        -  η·_d_0, π_1:h-1, ^*_1:h-1[(π_h(·|s_h)||π̂_h(·|s_h))] 
        = _d_0, π_1:h, ^*_1:h[V̂_h+1(s_h+1) - Q̂_h(s_h, a_h)] - _d_0, π̂_1:h, ^*_1:h[V̂_h+1(s_h+1) - Q̂_h(s_h, a_h)]
        -  η·_d_0, π_1:h-1, ^*_1:h-1[(π_h(·|s_h)||π̂_h(·|s_h))] 
        + _d_0, π_1:h, ^*_1:h[V^π_h+1(s_h+1)  - V̂_h+1(s_h+1)] - _d_0, π_1:h, ^*_1:h[V^π̂_h+1(s_h+1)  - V̂_h+1(s_h+1)].

The proposition can be obtained by iteratively using the above relationship for h∈ [H].




 §.§ Proof of Theorem <ref>





















First, with the assumption u^*∈ and ^*∈, the following lemma demonstrates that _t and _t are valid confidence sets.

    There exists an absolute constant c_1 such that for any δ∈ (0, 1], with probability at least 1-δ, for all t∈ [T], û∈, and ∈, it holds that
    
    L_t(û) - L_t(u^*) ≤ c_1 log(||T/δ),      L_t() - L_t(^*) ≤ c_1 log(||T/δ),

    which implies that u^*∈_t and ^*∈_t.








Then, we provide an additional lemma demonstrating the in-sample error of the MLE and optimistic estimators.

    There exists an absolute constant c_2 such that for any δ∈ (0, 1], with probability at least 1-δ, for all t∈ [T], we have
    
    ∑_i < t|σ(û_t(s^2_i,H, a^2_i,H) - û_t(s^1_i,H, a^1_i,H)) - σ(u^*(s^2_i,H, a^2_i,H) - u^*(s^1_i,H, a^1_i,H))|^2 ≤ c_2log(||T/δ);
    ∑_i < t|σ(ũ_t(s^2_i,H, a^2_i,H) - ũ_t(s^1_i,H, a^1_i,H)) - σ(u^*(s^2_i,H, a^2_i,H) - u^*(s^1_i,H, a^1_i,H))|^2 ≤ c_2log(||T/δ),

    and for all t∈ [T], h∈ [H], we have
    
    ∑_j∈{1,2}∑_h∈ [H]∑_i < t({d_0, π^j_i, [^*_1:h-1, _t,h,  ^*_h+1: H]}, {d_0, π^j_i, ^*_1:H})^2 ≤ c_2 log(||HT/δ);
    ∑_j∈{1,2}∑_h∈ [H]∑_i < t({d_0, π^j_i, [^*_1:h-1, _t,h,  ^*_h+1: H]}, {d_0, π^j_i, ^*_1:H})^2 ≤ c_2 log(||HT/δ),

    where ({d_0, π, }, {d_0, π', '}) denotes the TV distance between the probability distributions over the trajectories induced by d_0, π, and d_0, π', '.


    First, for ũ_t, we can obtain that with probability at least 1-δ, there exists an absolute constant c such that for all t∈ [T],
    
    ∑_i < t|σ(ũ_t(s^2_i,H, a^2_i,H) - ũ_t(s^1_i,H, a^1_i,H)) - σ(u^*(s^2_i,H, a^2_i,H) - u^*(s^1_i,H, a^1_i,H))|^2 
       ≤ c(∑_i< tlogz_i ·σ(u^*(s^1_i,H, a^1_i,H) - u^*(s^2_i,H, a^2_i,H))+ (1-z_i) ·σ(u^*(s^2_i,H, a^2_i,H) - u^*(s^1_i,H, a^1_i,H))/z_i ·σ(ũ_t(s^1_i,H, a^1_i,H) - ũ_t(s^2_i,H, a^2_i,H)) + (1- z_i) ·σ(ũ_t(s^2_i,H, a^2_i,H) - ũ_t(s^1_i,H, a^1_i,H)) + log(||T/δ))
        = c(L_t(u^*) - L_t(ũ_t) + log(||T/δ))
       ≤ c(L_t(u^*) - L_t(û_t) +  c_1 log(||T/δ) + log(||T/δ))
       ≤ c_2  log(||T/δ).

    where the first inequality is from Proposition B.2 from <cit.> and the second inequality uses Lemma <ref>. The result for û_t can be similarly established.

    Then, following similar steps, for _t,  we can obtain that with probability at least 1-δ, there exists an absolute constant c such that for all t∈  [T],
    
    ∑_j∈{1,2}∑_h∈ [H]∑_i < t({d_0, π^j_i, [^*_1:h-1, _t,h,  ^*_h+1: H]}, {d_0, π^j_i, ^*_1:H})^2
       ≤∑_j∈{1,2}∑_h∈ [H] c·(∑_i<tlog^*_h(s^j_i,h+1|s^j_i,h, a^j_i,h)/_t,h(s^j_i,h+1|s^j_i,h, a^j_i,h) + log(|_h|HT/δ))
        = c·(∑_j∈{1,2}∑_i<tlog^*,π^j_i(τ^j_i)/^π^j_i_t(τ^j_i) + 2log(||HT/δ))
        = c·(L_t(^*) - L_t(_t) + 2log(||HT/δ))
       ≤ c·(L_t(^*) - L_t(_t) + c_1 log(||T/δ) + 2log(||HT/δ))
       ≤ c_2 log(||HT/δ).

    The result for _t can also be similarly established.


 
In the following proofs, we omit the KL term in the decomposition to ease the presentation. Then, with probability at least 1-δ, for all t∈ [T], we can obtain that

    J(π^*) - J(π^1_t) 
        = _d_0, π^*, ^*[u^*(s_H, a_H)] - _d_0, π^1_t, ^*[u^*(s_H, a_H)] - (_d_0, π^*, ^*[û_t(s_H, a_H)] - _d_0, π^1_t, ^*[û_t(s_H, a_H)]) 
       + ∑_h∈ [H]_d_0, π^*, ^*[V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]- ∑_h∈ [H]_d_0, π^1_t, ^*[ V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]
       ≤_d_0, π_t^2, _t[ũ_t(s_H, a_H)] - _d_0, π^1_t, _t[ũ_t(s_H, a_H)] - (_d_0, π_t^2, _t[û_t(s_H, a_H)] - _d_0, π^1_t, _t[û_t(s_H, a_H)])_term (I)_t
       + ∑_h∈ [H]_d_0, π_t^2, _t[V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]+ ∑_h∈ [H]_d_0, π^1_t, ^*[ [_t, hV̂_t, h+1](s_h,a_h) - V̂_t, h+1(s_h+1)]_term (II)_t,

where the inequality is from the definition of π^2_t and the fact that (u^*, ^*)∈_t×_t from Lemma <ref>.

We define the following terms:

    term (A)_t    :=  _d_0, π_t^2, ^*[ũ_t(s_H, a_H)] - _d_0, π^1_t, ^*[ũ_t(s_H, a_H)] - (_d_0, π_t^2, ^*[u^*(s_H, a_H)] - _d_0, π^1_t, ^*[u^*(s_H, a_H)]),
    term (B)_t    := _d_0, π_t^2, ^*[u^*(s_H, a_H)] - _d_0, π^1_t, ^*[u^*(s_H, a_H)] - (_d_0, π_t^2, ^*[û_t(s_H, a_H)] - _d_0, π^1_t, ^*[û_t(s_H, a_H)]),
    term (C)_t    := ∑_j∈{1,2}∑_h∈ [H]_d_0, π^j_t, ^*[(_t,h(·|s_h, a_h), ^*_h(·|s_h, a_h))],
    term (D)_t    := ∑_j∈{1,2}∑_h∈ [H]_d_0, π^j_t, ^*[ (_t,h(·|s_h, a_h), ^*_h(·|s_h, a_h))].


For term (I)_t, we have that

    term (I)_t    := _d_0, π_t^2, _t[ũ_t(s_H, a_H)] - _d_0, π^1_t, _t[ũ_t(s_H, a_H)] - (_d_0, π_t^2, _t[û_t(s_H, a_H)] - _d_0, π^1_t, _t[û_t(s_H, a_H)])
        = _d_0, π_t^2, ^*[ũ_t(s_H, a_H)] - _d_0, π_t^1, ^*[ũ_t(s_H, a_H)] - (_d_0, π_t^2, ^*[u^*_t(s_H, a_H)] - _d_0, π_t^1, ^*[u^*_t(s_H, a_H)])
        + _d_0, π_t^2, ^*[u^*_t(s_H, a_H)] - _d_0, π_t^1, ^*[u^*_t(s_H, a_H)] - (_d_0, π_t^2, ^*[û_t(s_H, a_H)] - _d_0, π^1_t, ^*[û_t(s_H, a_H)]) 
        + _d_0, π_t^2, _t[ũ_t(s_H, a_H)] - _d_0, π_t^1, _t[ũ_t(s_H, a_H)] - (_d_0, π_t^2, ^*[ũ_t(s_H, a_H)] - _d_0, π_t^1, ^*[ũ_t(s_H, a_H)])
        + _d_0, π_t^2, ^*[û_t(s_H, a_H)] - _d_0, π^1_t, ^*[û_t(s_H, a_H)] - (_d_0, π_t^2, _t[û_t(s_H, a_H)] - _d_0, π^1_t, _t[û_t(s_H, a_H)])
       ≤_d_0, π_t^2, ^*[ũ_t(s_H, a_H)] - _d_0, π_t^1, ^*[ũ_t(s_H, a_H)] - (_d_0, π_t^2, ^*[u^*_t(s_H, a_H)] - _d_0, π_t^1, ^*[u^*_t(s_H, a_H)])
        + _d_0, π_t^2, ^*[u^*_t(s_H, a_H)] - _d_0, π_t^1, ^*[u^*_t(s_H, a_H)] - (_d_0, π_t^2, ^*[û_t(s_H, a_H)] - _d_0, π^1_t, ^*[û_t(s_H, a_H)])
        + 4B·({d_0, π^1_t, _t}, {d_0, π^1_t, ^*})+ 4B·({d_0, π^2_t, _t}, {d_0, π^2_t, })
       ≤_d_0, π_t^2, ^*[ũ_t(s_H, a_H)] - _d_0, π_t^1, ^*[ũ_t(s_H, a_H)] - (_d_0, π_t^2, ^*[u^*_t(s_H, a_H)] - _d_0, π_t^1, ^*[u^*_t(s_H, a_H)])_term (A)_t
        + _d_0, π_t^2, ^*[u^*_t(s_H, a_H)] - _d_0, π_t^1, ^*[u^*_t(s_H, a_H)] - (_d_0, π_t^2, ^*[û_t(s_H, a_H)] - _d_0, π^1_t, ^*[û_t(s_H, a_H)])_term (B)_t
        + 4B·∑_j∈{1,2}∑_h∈ [H]_d_0_π^j_t, ^*[(_t,h(·|s_h, a_h), ^*_h(·|s_h, a_h))]_term (C)_t.


For term (II)_t, we have that

    term (II)_t    = ∑_h∈ [H]_d_0, π_t^2, _t[V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]
        + ∑_h∈ [H]_d_0, π^1_t, ^*[ [_t, hV̂_t, h+1](s_h,a_h) - V̂_t, h+1(s_h+1)]
        = ∑_h∈ [H]_d_0, π_t^2, ^*[V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]
        + ∑_h∈ [H]_d_0, π_t^2, _t[V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]
        - ∑_h∈ [H]_d_0, π_t^2, ^*[V̂_t, h+1(s_h+1) - [_t, hV̂_t, h+1](s_h,a_h)]
        + ∑_h∈ [H]_d_0, π^1_t, ^*[ [_t, hV̂_t, h+1](s_h,a_h) - V̂_t, h+1(s_h+1)]
       ≤ 2B·∑_j∈{1,2}∑_h∈ [H]_d_0, π_t^j, ^*[(_t,h(·|s_h, a_h)), ^*_h(·|s_h, a_h)]
        + 2BH ·({d_0, π_t^2, _t}, {d_0, π_t^2, ^*})
       ≤ 2B·∑_j∈{1,2}∑_h∈ [H]_d_0, π_t^j, ^*[(_t,h(·|s_h, a_h)), ^*_h(·|s_h, a_h)]_term (D)_t
        + 2BH ·∑_j∈{1,2}∑_h∈ [H]_d_0, π_t^j, ^*[(_t,h(·|s_h, a_h)), ^*_h(·|s_h, a_h)]_term (C)_t.


In the above derivations, we have repeatedly used similar relationships as follows:

    ({d_0, π^2_t, _t}, {d_0, π^2_t, ^*}) ≤∑_h∈ [H]_d_0, π^2_t, ^*[(_t,h(·|s_h, a_h), ^*_h(·|s_h, a_h))],

which can be derived as

    ({d_0, π^2_t, _t}, {d_0, π^2_t, ^*})    ≤∑_h∈ [H]({d_0, π^2_t, ^*_1:h-1, _t, h:H}, {d_0, π^2_t, ^*_1:h, _t, h+1:H})
       = ∑_h∈ [H]_d_0, π^2_t, ^*[(_t, h(·|s_h,a_h), ^*_h(·|s_h,a_h)})].


Then, we can obtain that

    ∑_t∈ [T]J(π^*) - J(π̂^1_t)    ≤∑_t∈ [T]term (A)_t + ∑_t∈ [T]term (B)_t + (4B + 2BH) ∑_t∈ [T]term (C)_t + 2B ∑_t∈ [T]term (D)_t.


Then, we control the sum of each individual term in the following. First, for term (A)_t, with probability at least 1-δ, we have that

    ∑_t∈ [T]term (A)_t 
       = ∑_t∈ [T]_d_0, π_t^2, ^*[ũ_t(s_H, a_H)] - _d_0, π^1_t, ^*[ũ_t(s_H, a_H)] - (_d_0, π_t^2, ^*[u^*(s_H, a_H)] - _d_0, π^1_t, ^*[u^*(s_H, a_H)])
       ≤∑_t∈ [T]ũ_t(s^2_t,H, a^2_t,H) - ũ_t(s^1_t,H, a^1_t,H) - (u^*(s^2_t,H, a^2_t,H) - u^*(s^1_t,H, a^1_t,H)) + O(B√(Tlog(1/δ)))
       ≤√(d_∑_t=2^T(1+ ∑_i=1^t-1(ũ_t(s^2_i,H, a^2_i,H) - ũ_t(s^1_i,H, a^1_i,H) - (u^*(s^2_i,H, a^2_i,H) - u^*(s^1_i,H, a^1_i,H)))^2)) + O(B√(Tlog(1/δ)))
       ≤√(d_∑_t=2^T(1+ κ^-2∑_i=1^t-1(σ(ũ_t(s^2_i,H, a^2_i,H) - ũ_t(s^1_i,H, a^1_i,H)) - σ(u^*(s^2_i,H, a^2_i,H) - u^*(s^1_i,H, a^1_i,H)))^2)) + O(B√(Tlog(1/δ)))
       ≲κ^-1 B √(d_ T log(||T/δ)),

where the first inequality is from the Hoeffding inequality, the second inequality uses the Eluder coefficient d_ := (1,  - , T) from Definition <ref>, the third inequality leverages the mean value theorem with κ:= 1/(2+ exp(-B)+ exp(B)) representing the minimum derivative of σ(·) in the regime of [0, B], and the last inequality incorporates Lemma <ref>. A similar result can be obtained for term (B)_t.

For term (C)_t, we have that

    ∑_t∈ [T]term (C)_t    = ∑_j∈{1,2}∑_t∈ [T]∑_h∈ [H]_d_0, π^j_t, ^*[(_t,h(·|s_h, a_h), ^*_h(·|s_h, a_h))]
        = ∑_j∈{1,2}∑_t∈ [T]∑_h∈ [H]({d_0, π^j_t, [^*_1:h-1, _t,h,  ^*_h+1: H]}, {d_0, π^j_t, ^*_1:H})
       ≤ 2H·ξ(d_, T, c_2log(||HT/δ)),

where the last step is from the generalized Eluder-type condition in Definition <ref> and Lemma <ref>. A similar result can be obtained for term (D)_t. 

Finally, we obtain that

    Reg(T) ≲   κ^-1B√(d_Tlog(||T/δ)) + B^2Hξ(d_, T, c_2log(||HT/δ) 
       - η·∑_h∈ [H]_d_0, π^*, ^*[(π^*_h(·|s_h), π^1_t, h(·|s_h))],

which concludes the proof.






§ TECHNICAL LEMMAS


  Given a loss functional with respect to p(·|x), written as 
    _w ∼ p(·)[-U(w) + η(p(·),p_0(·) )] = η( p(·), p_0(·)exp(1/ηU(·)) )- η·log_w ∼ p_0(·)exp(1/ηU(w))_C_r,

where the minimizer of the loss functional is 
p^*(w) = 1/C_r p_0(w)exp(1/η U(w)), also known as Gibbs distribution.



    Given a function class , its Eluder coefficient (λ, , T) is defined as the smallest number d so that for any sequence {x_t: t∈ [T]} and {f_t: t∈ [T]}∈, 
    
    ∑_t = 2^T |f_t(x_t) - f^*(x_t)| ≤√(d∑_t = 2^T(λ + ∑_i = 1^t-1 (f_t(x_i) - f^*(x_i))^2)).




    There exists a real number d_∈^+ and a function ξ such that for any (T, Δ) ∈×^+, transitions {'_t: t∈ [T]} and policies {π_t: t∈ [T]}, we have
    
    ∀ t∈ [T],   ∑_i < t({d_0, '_i, π_i}, {d_0, , π_i})^2  ≤Δ  ⇒∑_t ∈ [T]({d_0, '_t, π_t}, {d_0, , π_t}) ≤ξ(d_, T, Δ).



