Building Math Agents with Multi-Turn Iterative Preference Learning

Abstract

Recent studies have shown that large language models’ (LLMs) mathematical problem-solving capabilities can be enhanced by integrating external tools, such as code interpreters, and employing multi-turn Chain-of-Thought (CoT) reasoning. While current methods focus on synthetic data generation and Supervised Fine-Tuning (SFT), this paper studies the complementary direct preference learning approach to further improve model performance. However, existing direct preference learning algorithms are originally designed for the single-turn chat task, and do not fully address the complexities of multi-turn reasoning and external tool integration required for tool-integrated mathematical reasoning tasks. To fill in this gap, we introduce a multi-turn direct preference learning framework, tailored for this context, that leverages feedback from code interpreters and optimizes trajectory-level preferences. This framework includes multi-turn DPO and multi-turn KTO as specific implementations. The effectiveness of our framework is validated through training of various language models using an augmented prompt set from the GSM8K and MATH datasets. Our results demonstrate substantial improvements: a supervised fine-tuned Gemma-1.1-it-7B model’s performance 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.

1 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 (OpenAI, 2023), Claude (Anthropic, 2023), and Gemini (Team et al., 2023). 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 (Hendrycks et al., 2021; Cobbe et al., 2021b; Zheng et al., 2021). 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 (Cobbe et al., 2021b; Shao et al., 2022; Mishra et al., 2022; Zhang et al., 2024a). 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 (Gou et al., 2023b; Toshniwal et al., 2024; Shao et al., 2024).

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 (Gou et al., 2023b), MetaMathQA (Yu et al., 2023), MAmmoTH (Yue et al., 2023, 2024), and Open-MathInstruct (Toshniwal et al., 2024). These methods and synthetic datasets have yielded significant improvements in test accuracy on standard benchmarks like MATH (Hendrycks et al., 2021) and GSM8K (Cobbe et al., 2021a).

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 (Bai et al., 2022; Ouyang et al., 2022; Touvron et al., 2023; Team et al., 2023). Broadly speaking, the RLHF learning paradigm, which was originally designed for aligning large language models (LLMs) with human values and preferences (Bai et al., 2022; Ouyang et al., 2022), is distinct from SFT as it learns from relative feedback (Christiano et al., 2017; Ziegler et al., 2019). 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 (Bai et al., 2022). 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 (Schulman et al., 2017)) are often sample inefficient and unstable (Choshen et al., 2019), our goal is to derive direct preference learning algorithms that directly learn from the preference dataset (Zhao et al., 2023; Rafailov et al., 2023; Azar et al., 2023).

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 (Xiong et al., ; Rafailov et al., 2023). 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 (Xiong et al., ; Guo et al., 2024b). Finally, we evaluate our approach through case studies using augmented training sets from MATH and GSM8K benchmarks, employing various base models such as Gemma (Team et al., 2024), CodeGemma (Team, 2024), and Mistral (Jiang et al., 2023). 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.

1.1 Problem Formulation

We denote prompt as $x\in 𝒳$ 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\in [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 ${\pi }_h(\cdot |s_h)\in \mathrm{\Delta }(𝒜)$. 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) (Yao et al., 2022).

• 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,\mathrm{\cdots },a_h,o_h),$$

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

$$\tau =(x,a_1,o_1,\mathrm{\cdots },o_{H-1},a_H).$$

See Figure 1 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 Xiong et al. . Formally, we define the following MDP.

Definition 1.

An MDP is specified by a tuple $(𝒮,𝒜,H,{\mathbb{P}}^{\ast },d_0)$, where $𝒜$ is the action space, $H$ is the episode length, ${\mathbb{P}}^{\ast }={\{{\mathbb{P}}_h^{\ast }\}}_{h=1}^H$ are the state transition kernels, and $d_0$ denotes the distribution of prompt $s_1=x$. For each $h\in [H]$, ${\mathbb{P}}_h^{\ast }(\cdot |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 $\tau =(x,a_1,o_1,\mathrm{\cdots },o_{H-1},a_H)$ is generated by: $s_1=x\sim d_0$ and for all $h\in [H],a_h\sim {\pi }_h(\cdot |s_h),o_h\sim {\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)$ where $s_{h+1}=(s_h,a_h,o_h)$. When there is no ambiguity, the abbreviation $s_{h+1}\sim {\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)$ is also adopted.

Figure 1: An example of multi-turn mathematical reasoning with Python interpreter. The action is in a ReAct style (Yao et al., 2022) where it consists of a reasoning step $f_h$ and an executation step $e_h$.

The MDP formulation of preference learning was recently studied in Zhong et al. (2024); Rafailov et al. (2024); Xie et al. (2024a) 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 Ziegler et al. (2019); Ouyang et al. (2022); Bai et al. (2022) to assume that the preference signal is generated by the so-called Bradley-Terry model.

Definition 2 (Bradley-Terry model).

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

$$\text{Prob}\left(y^1\succ y^2\mid x,y^1,y^2\right)=\sigma \left(u^{\ast }(x,y^1)-u^{\ast }(x,y^2)\right),$$ (1)

where $\sigma$ is the sigmoid function $\sigma (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\succ y^2$ while $z=0$ indicating $y^2\succ 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 (Christiano et al., 2017; Ziegler et al., 2019), 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.

Example 1 (Result Checking in Math).

Since the math reasoning datasets GSM8K (Cobbe et al., 2021a) and MATH (Hendrycks et al., 2021) have the gold answer, we can check the final answer to determine the reward. In this case, $u^{\ast }(x,y)=𝕀(a_H=\text{gold answer})$.

Example 2 (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 Cobbe et al. (2021b); Lightman et al. (2023), 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.

Example 3 (Process-supervised Reward Model (PRM) and PRM without Human Annotation.).

Lightman et al. (2023) 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. Wang et al. (2023a) 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):={\displaystyle \frac{{\sum }_{j=1}^N𝕀(a_H^j=\text{gold answer})}{N}}.$ (2)

We may also use a hard version

$r(s_h,a_h):=𝕀(\text{There exists a }{j}_0:a_H^{j_0}=\text{gold answer}).$ (3)

Then, we can train the PRM by

$${ℒ}_{PRM}(\theta )={𝔼}_{\tau \sim 𝒟}\left[{\sum }_{h=1}^{H}r(s_h,a_h)\log r_{\theta }+(1-r(s_h,a_h))\log (1-r_{\theta })\right].$$ (4)

In this case, we can use $u^{\ast }(x,y)={\min }_{h\in [H]}{r}_{\theta }(s_h,a_h)$ (Lightman et al., 2023), where $r_{\theta }$ is the constructed step-wise reward function.

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

1.2 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 (Wei et al., 2022; Zhou et al., 2022; Zhu et al., 2022; Tong et al., 2024), 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 (Cobbe et al., 2021a; Hendrycks et al., 2021). To overcome these limitations, several studies have explored the use of external tools to enhance the LLMs’ problem-solving abilities. This includes calculators (Cobbe et al., 2021b; Shao et al., 2022), symbolic solvers (Zhang, 2023), and code interpreters (Mishra et al., 2022; OpenAI, 2023). 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 (Gao et al., 2023a; Chen et al., 2022). 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 (Gou et al., 2023a). In view of this, tool-integrated reasoning is proposed to combine the natural-language-based intrinsic reasoning with the external tools (Gou et al., 2023b) and has achieved great progresses in recent studies (Gou et al., 2023b; Yue et al., 2023; Yu et al., 2023; Shao et al., 2024; Toshniwal et al., 2024). 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) (Schulman et al., 2017), which leads to the great successes in Chat-GPT (OpenAI, 2023), Gemini (Team et al., 2023), and Claude (Anthropic, 2023). However, applying PPO requires extensive efforts and resources (Choshen et al., 2019; Engstrom et al., 2020), 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 (Dong et al., 2023) and was later extended to machine translation (Gulcehre et al., 2023) and mathematical problem solving (Yuan et al., 2023a). Its theoretical advantage was explored in Gui et al. (2024). Subsequently, another long line of works proposes direct preference learning algorithms, including SLiC (Zhao et al., 2023), DPO (Rafailov et al., 2023), IPO (Azar et al., 2023), KTO (Ethayarajh et al., 2024), and GPO (Tang et al., 2024). 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 Munos et al. (2023); Swamy et al. (2024); Ye et al. (2024); Rosset et al. (2024) beyond the reward-based framework or post-processing of the model (Lin et al., 2023; Zheng et al., 2024).

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 (Tunstall et al., 2023). After this, a long line of work (e.g., Liu et al., 2023b; Xiong et al., ; Guo et al., 2024b; Xu et al., 2023; Tajwar et al., 2024; Xie et al., 2024a; Zhang et al., 2024b; Liu et al., 2024a, b; Meng et al., 2024) 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 (Xiong et al., ; Xu et al., 2023; Hoang Tran, 2024) and its variants (e.g., Chen et al., 2024b; Rosset et al., 2024; Cen et al., 2024; Zhang et al., 2024c) have made state-of-the-art open-source models (Dong et al., 2024), or even the industry models (qwe, 2024; Meta, 2024). 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 Shani et al. (2024), 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) (Dong et al., 2023; Yuan et al., 2023b; Touvron et al., 2023) is extensively employed for synthetic data generation, whether through on-policy (self-improving) (Yuan et al., 2023a) or off-policy (knowledge distillation) methods (Gou et al., 2023b; Yu et al., 2023; Toshniwal et al., 2024; Singh et al., 2023; Tong et al., 2024). The reward signal in these scenarios is typically derived from either final result checking or Outcome-supervised Reward Models (ORMs) (Uesato et al., 2022; Zelikman et al., 2022). A novel approach by Lightman et al. (2023) 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 (Lightman et al., 2023; Wang et al., 2023a).

In addition to the RAFT, the GRPO algorithm proposed in Shao et al. (2024) 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 (Williams, 1992) thus falling into the scope of deep RL methods.

Further advancements include adapting direct preference learning algorithms to mathematical problem solving. For instance, Jiao et al. (2024); Yuan et al. (2024) have applied the original DPO or KTO by taking the trajectory completion as a “meta” action. Xie et al. (2024b); Pang et al. (2024) further adapt the online iterative DPO originally designed for chat (Xiong et al., ; Xu et al., 2023; Hoang Tran, 2024) 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, Xie et al. (2024b); Chen et al. (2024a); Lai et al. (2024) leverage Monte Carlo Tree Search (MCTS) and use the estimated Q value to generate the proxy labels for the intermediate steps. Lai et al. (2024) proposes to use AI feedback like GPT-4 (Lai et al., 2024) to find the first error step in the trajectory. Meanwhile, Lu et al. (2024) 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 (Pi et al., 2024). Finally, a very recent study by Chen et al. (2024a) 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 1.1 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 (11)). 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.

2 Algorithms Development

We develop the main algorithms of this paper in this section. We proceed to handle the general MDP formulation presented in Section 1.1, 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..

2.1 Planning with a Fixed Model: Optimality Condition

Following Rafailov et al. (2023), we first establish the connection between any model $ℳ=(𝒮,𝒜,H,\mathbb{P},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 ${\pi }_{\mathrm{ref}}$:

$$\underset{\pi }{argmax}J(\pi ;ℳ,{\pi }_{\mathrm{ref}})={𝔼}_{x\sim d_0}{𝔼}_{a_h\sim {\pi }_h(\cdot |s_h),o_h\sim {\mathbb{P}}_h(\cdot |s_h,a_h)}[u(x,y)-\eta \sum _{h=1}^H{D}_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))].$$ (5)

In the single-turn case (i.e., $H=1$ and without transitions $\mathbb{P}$), Rafailov et al. (2023); Azar et al. (2023) show that the optimal solution with respect to a utility function $u$ admits a closed-form solution, which is the Gibbs distribution (see Lemma 3):

$${\pi }_{ℳ}(a_1|x)\propto {\pi }_{\mathrm{ref}}(a_1|x)\exp \left(\frac{u(x,a_1)}{\eta }\right).$$

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 (Williams and Peng, 1991; Ziebart, 2010).

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\sim d_0}[{𝔼}_{a_1\sim {\pi }_1(\cdot |x)}\left[{𝔼}_{o_1\sim {\mathbb{P}}_1(\cdot |x,a_1)}\underset{\text{Inner Loop}}{\underbrace{{𝔼}_{a_2\sim {\pi }_2(\cdot |s_2)}u(s_2,a_2)-\eta D_{\mathrm{KL}}({\pi }_2(\cdot |s_2),{\pi }_{\mathrm{ref},2}(\cdot |s_2))}}\right]-\eta D_{\mathrm{KL}}({\pi }_1(\cdot |s_1),{\pi }_{\mathrm{ref},1}(\cdot |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 3 to solve

$${\pi }_{ℳ,2}(\cdot |s_2)=\underset{{\pi }_2}{argmax}{𝔼}_{a_2\sim {\pi }_2(\cdot |s_2)}(u(s_2,a_2)-\eta \cdot D_{\mathrm{KL}}({\pi }_2(\cdot |s_2),{\pi }_{\mathrm{ref},2}(\cdot |s_2)))\propto {\pi }_{\mathrm{ref},2}(\cdot |s_2)\cdot \exp \left(\frac{u(s_2,\cdot )}{\eta }\right).$$

Then, we can define the value of the inner loop associated with ${\pi }_{ℳ,2}$ as

$V_{ℳ,2}(s_2)$ $:={𝔼}_{a_2\sim {\pi }_{ℳ,2}(\cdot |s_2)}[u(s_2,a_2)-\eta D_{\mathrm{KL}}({\pi }_{ℳ,2}(\cdot |s_2),{\pi }_{\mathrm{ref},2}(\cdot |s_2))]$

$Q_{ℳ,1}(s_1,a_1)$ $:={𝔼}_{o_1\sim {\mathbb{P}}_1(\cdot |s_1,a_1)}\left[V_{ℳ,2}(s_2)\right].$

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

${\pi }_{ℳ,1}(\cdot |s_1)=\underset{{\pi }_1}{argmax}{𝔼}_{a_1\sim {\pi }_1(\cdot |x)}[Q_{ℳ,1}(s_1,a_1)-\eta D_{\mathrm{KL}}({\pi }_1(\cdot |s_1),{\pi }_{\mathrm{ref},1}(\cdot |s_1))]\propto {\pi }_{\mathrm{ref},1}(\cdot |s_1)\cdot \exp \left({\displaystyle \frac{Q_{ℳ,1}(s_1,\cdot )}{\eta }}\right).$

By construction, it can be observed that ${\{{\pi }_{ℳ,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

(6)

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

${\pi }_{ℳ,h}(a_h|s_h)$ $:={\displaystyle \frac{1}{Z_h(s_h)}}{\pi }_{\mathrm{ref},h}(a_h|s_h)\cdot \exp \left({\displaystyle \frac{Q_{ℳ,h}(s_h,a_h)}{\eta }}\right)\text{(Gibbs distribution of }{Q}_{ℳ,h}\text{)}$ (7)

$V_{ℳ,h}(s_h)$ $:={𝔼}_{a_h\sim {\pi }_{ℳ,h}(\cdot |s_h)}[Q_{ℳ,h}(s_h,a_h)-\eta \cdot D_{\mathrm{KL}}({\pi }_{ℳ,h}(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))]$

$=\eta \log {𝔼}_{{\pi }_{\mathrm{ref},h}(a_h^{\prime }|s_h)}\exp \left({\displaystyle \frac{Q_{ℳ,h}(s_h,a_h^{\prime })}{\eta }}\right),$

where $Z_h(s_h)={\sum }_{a_h\in 𝒜}{\pi }_{\mathrm{ref},h}(a_h|s_h)\cdot \exp \left(\frac{Q_{ℳ,h}(s_h,a_h)}{\eta }\right)$ is the normalization constant. The second equality in the definition of the $V$-value is from Lemma 3. Then, by definition, ${[{\pi }_{ℳ,h}]}_{h=1}^H$ is the optimal policy. Essentially, we solve $H$ Gibbs distributions in terms of the $Q$-values.

2.2 Planning with a Fixed Model: Practical Algorithm

While (7) can be approximately solved with standard deep RL methods, here we are interested in the implementation in a direct preference learning manner like SLiC (Zhao et al., 2023), DPO (Rafailov et al., 2023) or IPO (Azar et al., 2023). The existing attempts (e.g., Yuan et al., 2024) 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,{\tau }^w,{\tau }^l)$, where ${\tau }^w$ refers to the preferred (i.e., winning) trajectory, we have

$-\log \sigma \left(\eta \left[\log {\displaystyle \frac{{\text{Prob}}_{\pi }({\tau }^l|x)}{{\text{Prob}}_{{\pi }_{\mathrm{ref}}}({\tau }^l|x)}}-\log {\displaystyle \frac{{\text{Prob}}_{\pi }({\tau }^w|x)}{{\text{Prob}}_{{\pi }_{\mathrm{ref}}}({\tau }^w|x)}}\right]\right)$ (8)

$=-\log \sigma \left(\eta \left[\log {\displaystyle \prod _{h=1}^H}{\displaystyle \frac{{\pi }_h(a_h^l|s_h^l)\cancel{{\mathbb{P}}_h(o_h^l|s_h^l,a_h^l)}}{{\pi }_{\mathrm{ref},h}(a_h^l|s_h^l)\cancel{{\mathbb{P}}_h(o_h^l|s_h^l,a_h^l)}}}-\log {\displaystyle \prod _{h=1}^H}{\displaystyle \frac{{\pi }_h(a_h^w|s_h^w)\cancel{{\mathbb{P}}_h(o_h^w|s_h^w,a_h^w)}}{{\pi }_{\mathrm{ref},h}(a_h^w|s_h^w)\cancel{{\mathbb{P}}_h(o_h^w|s_h^w,a_h^w)}}}\right]\right)$

$=-\log \sigma \left(\eta {\displaystyle \sum _{h=1}^H}\left[\log {\displaystyle \frac{{\pi }_h(a_h^l|s_h^l)}{{\pi }_{\mathrm{ref},h}(a_h^l|s_h^l)}}-\log {\displaystyle \frac{{\pi }_h(a_h^w|s_h^w)}{{\pi }_{\mathrm{ref},h}(a_h^w|s_h^w)}}\right]\right).$

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 {\displaystyle \frac{{\pi }_{ℳ,h}(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}+\eta \log {𝔼}_{{\pi }_{\mathrm{ref},h}(a_h^{\prime }|s_h)}\exp \left({\displaystyle \frac{Q_{ℳ,h}(s_h,a_h^{\prime })}{\eta }}\right)$ (9)

$=\log {\displaystyle \frac{{\pi }_{ℳ,h}(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}+V_{ℳ,h}(s_h),$

where two equalities uses the definition of the optimal policy ${\pi }_{ℳ,h}$ and $V$-value $V_{ℳ,h}$ in (7), respectively. Furthermore, by the definition of $Q$-values $Q_{ℳ,h}$ in (6), we have

${𝔼}_{o_h\sim {\mathbb{P}}_h(\cdot |s_h,a_h)}{V}_{ℳ,h+1}(s_{h+1})$ $=\log {\displaystyle \frac{{\pi }_{ℳ,h}(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}+V_{ℳ,h}(s_h),\text{if }h\le H-1$ (10)

$u(s_H,a_H)$ $=\log {\displaystyle \frac{{\pi }_{ℳ,H}(a_H|s_H)}{{\pi }_{\mathrm{ref},H}(a_H|s_H)}}+V_{ℳ,H}(s_H).$

Summing over $h\in [H]$, we have

$u(s_H,a_H)$ $=\eta {\displaystyle \sum _{h=1}^H}\log {\displaystyle \frac{{\pi }_{ℳ,h}(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}+{\displaystyle \sum _{h=1}^H}\left[V_{ℳ,h}(s_h)-{𝔼}_{o_h\sim {\mathbb{P}}_h(\cdot |s_h,a_h)}{V}_{ℳ,h+1}(s_{h+1})\right]$ (11)

$=\underset{\text{term }(A)}{\underbrace{\eta {\displaystyle \sum _{h=1}^H}\log {\displaystyle \frac{{\pi }_{ℳ,h}(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}}}+\underset{\text{term }(B)}{\underbrace{V_{ℳ,1}(s_1)}}+\underset{\text{term }(C)}{\underbrace{{\displaystyle \sum _{h=1}^{H-1}}\left[V_{ℳ,h+1}(s_{h+1})-{𝔼}_{o_h\sim {\mathbb{P}}_h(\cdot |s_h,a_h)}{V}_{ℳ,h+1}(s_{h+1})\right]}}.$

Here, term $(A)$ is the counterpart of $\eta \log \frac{\pi (a_1|s_1)}{{\pi }_{\mathrm{ref}}(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|\le 4{\left[\sum _{h=1}^{H-1}{\sigma }_h^2\right]}^{1/2},$$

where ${\sigma }_h^2$ is the conditional variance of $V_{ℳ,h+1}(s_{h+1})-{𝔼}_{o_h\sim {\mathbb{P}}_h(\cdot |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 $\text{term }(C)=0$. Thus, with a dataset $𝒟$ consisting of $(x,{\tau }^w,{\tau }^l)$, the following multi-turn DPO (M-DPO) loss can be adopted:

$${ℒ}_{\text{M-DPO}}(\theta )=-\sum _{(x,{\tau }^w,{\tau }^l)\in 𝒟}\log \sigma \left(\eta \sum _{h=1}^H\left[\log \frac{{\pi }_{\theta ,h}(a_h^l|s_h^l)}{{\pi }_{\mathrm{ref},h}(a_h^l|s_h^l)}-\log \frac{{\pi }_{\theta ,h}(a_h^w|s_h^w)}{{\pi }_{\mathrm{ref},h}(a_h^w|s_h^w)}\right]\right),$$ (12)

We emphasize again that although the loss presented in (12) is identical to the one in (8), a rigorous derivation procedure (rather than a direct plug-in) is provided. To the best of our knowledge, (12) 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 (11) implying that with term $(C)=0$, the implicit reward is given by $A=\eta {\sum }_{h=1}^H\log \frac{{\pi }_h^{\ast }(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}$, a multi-turn version of KTO (Ethayarajh et al., 2024), denoted as M-KTO, can also be naturally derived:

$${ℒ}_{\text{M-KTO}}(\theta )={𝔼}_{x,y\sim 𝒟}\left[{\lambda }_y-v(x,y)\right],$$ (13)

where

$u_{\theta }(x,y)$ $=\eta {\displaystyle \sum _{h=1}^H}\log {\displaystyle \frac{{\pi }_{u,h}(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}},$

$z_0$ $={𝔼}_{x^{\prime }\sim 𝒟,{\tau }^{\prime }\sim {\pi }_{\theta }(\cdot |x^{\prime })}{\displaystyle \sum _{h=1}^H}{D}_{\mathrm{KL}}({\pi }_{\theta }(\cdot |s_h),{\pi }_{\mathrm{ref}}(\cdot |s_h)),$

and

Here ${\lambda }_{+}$ and ${\lambda }_{-}$ are two hyper-parameters. We notice that Mitra et al. (2024) 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 (12) and (13), 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 (Shani et al., 2024). In these scenarios, (12) is biased and cannot lead to the optimal policy since $\text{term }(C)\ne 0$. Instead, one should first construct a value network based on the Bellman equations provided in (6) and (7), similar to the approach in Richemond et al. (2024). Subsequently, $\text{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 Zhong et al. (2024); Rafailov et al. (2024); Xie et al. (2024a) 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}$.

2.3 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 (Xiong et al., ; Ye et al., 2024; Guo et al., 2024b; Rosset et al., 2024; Dong et al., 2024; Tajwar et al., 2024) and DRL-based approach or rejection sampling fine-tuning (Bai et al., 2022; Ouyang et al., 2022; Touvron et al., 2023). 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:

$$\underset{\pi }{\max }{𝔼}_{x\sim d_0}{𝔼}_{a_h\sim \pi (\cdot |s_h),o_h\sim {\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)}[u^{\ast }(x,y)-\eta \sum _{h=1}^H{D}_{\mathrm{KL}}(\pi (\cdot |s_h),{\pi }_0(\cdot |s_h))],$$ (14)

i.e., ${\max }_{\pi }J(\pi ;{ℳ}^{\ast },{\pi }_0)$ where ${ℳ}^{\ast }=(𝒮,𝒜,H,{\mathbb{P}}^{\ast },d_0,u^{\ast })$ is the groundtruth environment and ${\pi }_0$ is the initial policy (e.g., from SFT) that RLHF starts from. This target is widely adopted in practice (Christiano et al., 2017; Ouyang et al., 2022; Bai et al., 2022; Rafailov et al., 2023; Dong et al., 2024) and requires us to search for the optimal policy only at a fixed KL ball centered at the SFT policy ${\pi }_0$ (Xiong et al., ; Ye et al., 2024; Xie et al., 2024a).

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

$$\underset{\pi }{\max }{𝔼}_{x\sim d_0}{𝔼}_{a_h\sim \pi (\cdot |s_h),o_h\sim {\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)}\left[u^{\ast }(x,y)\right].$$ (15)

This target is the standard one in canonical RL studies (Sutton and Barto, 2018). 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 (14) is known to be beneficial for mitigating over-fitting issue and forgetting on the out-of-domain tasks (Gao et al., 2023b; Lin et al., 2023; Coste et al., 2023). On the other hand, (15) 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 (15) 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 1. 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 (7) that the optimal policy is a layer-wise Gibbs distribution, which generalizes the bandit formulation in Xiong et al. . Specifically, starting from ${\pi }_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 1. 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 ${\pi }_t^1$, where the planning algorithms have been discussed in Section 2.2, e.g., optimizing the M-DPO or M-KTO loss in (12) or (13). However, it is widely recognized in RL studies (Sutton and Barto, 2018; Auer et al., 2002) 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 ${\pi }_t^1$ given the historical information we collect so far into its policy choice. We call the policy of the second agent ${\pi }_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 (Anthropic, 2023), the authors choose to use the checkpoints from different training steps to collect data;

• Inference parameters tuning: in the LLaMA project (Touvron et al., 2023), the authors carefully tune the sampling temperature to balance data diversity and data quality;

• West-of-n sampling: Xu et al. (2023); Hoang Tran (2024); Pace et al. (2024); Dong et al. (2024) 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 (14) and (15) 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 (14): if we fix the reference model as the initial policy, i.e., ${\pi }_{t,\mathrm{ref}}={\pi }_0,\forall t\in [T]$, we always search the optimal policy within the KL ball centered at ${\pi }_0$, and thus optimize the KL-regularized target.

• Non-regularized target in (15): in contrast, inspired by the mirror descent (Nemirovskij and Yudin, 1983), if we update the reference policy every iteration to be the policy learned in the last iteration, i.e., ${\pi }_{t,\mathrm{ref}}={\pi }_{t-1}^1,\forall t\in [T]$, the cumulative update can make the model to move away from the original ${\pi }_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 2 to facilitate the understanding.

Figure 2: Illustration of the difference between the two learning objectives. The left-hand figure corresponds to the KL-regularized target where we do not update the reference model. The right-hand figure corresponds to the non-regularized target where we always update the reference model as the last-iteration one.

Algorithm 1 Online Iterative M-GSHF

1:  Input: KL coefficient $\eta >0$, horizon $T>0$, initial policy ${\pi }_0$, batch size $m>0$.

2:  Initialize $𝒟\leftarrow \mathrm{\varnothing }$ and ${\pi }_1^1={\pi }_1^2={\pi }_{1,\mathrm{ref}}\leftarrow {\pi }_0$.

3:  for $t=1,2,\mathrm{\cdots },T$ do

4:     Sample $m$ pairs $(x,{\tau }^1,{\tau }^2,z)$ as ${𝒟}_t$ by $x\sim d_0,{\tau }^1\sim {\pi }_t^1,{\tau }^2\sim {\pi }_t^2$, receive the $m$ preference signals $z$ following the Bradley-Terry model from Definition 1 and update the preference dataset $𝒟\leftarrow 𝒟\cup {𝒟}_t$.

5:     $▷$ Extract the empirically optimal policy from historical data

6:     Practical: Perform the planning algorithms on $𝒟$ to get ${\pi }_t^1$ (e.g., using the M-DPO loss in (12) or the M-KTO loss in (13))

7:     Theoretical: Perform MLE on $𝒟$ to obtain model estimation ${\widehat{ℳ}}_t=({\widehat{u}}_t,{\widehat{\mathbb{P}}}_t)$ as in (18) and (19); call Oracle 3 with ${\widehat{ℳ}}_t,\eta ,{\pi }_{t,\mathrm{ref}}$ to get ${\pi }_t^1$

8:     $▷$ Select the exploration policy to facilitate learning

9:     Practical: Given ${\pi }_t^1$, select ${\pi }_t^2$ as an exploration policy using heuristic methods (such as mixture sampling, inference parameters tuning and west-of-n sampling listed in Section 2.3)

10:     Theoretical: Given ${\pi }_t^1$, choose ${\pi }_t^2$ as an exploration policy following (20)

11:     $▷$ Choose the reference model to control regularization level

12:     if KL-regularized target in (14) then

13:        Keep ${\pi }_{t+1,\mathrm{ref}}\leftarrow {\pi }_0$

14:     else if Non-regularized target in (15) then

15:        Update ${\pi }_{t+1,\mathrm{ref}}\leftarrow {\pi }_t^1$

16:     end if

17:  end for

18:  Output: the best model in ${\pi }_{1:T}^1$ by a validation set.

2.4 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 Wang et al. (2023b) while the techniques of analyzing mirror-descent-style algorithm and corresponding guarantees have also be developed in Cai et al. (2020), 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

$${\pi }^{\ast }:=\underset{\pi }{argmax}J(\pi ):=J(\pi ;{ℳ}^{\ast },{\pi }_0),$$ (16)

and introduce the standard notion of regret as

$$\text{Reg}(T):=\sum _{t\in [T]}J({\pi }^{\ast })-J({\pi }_t^1),$$ (17)

which represents the cumulative performance loss over $T$ steps comparing the learned policies ${[{\pi }_t^1]}_{t=1}^T$ against the optimal policy ${\pi }^{\ast }$. In addition, we consider that a bounded $u^{\ast }(x,y)\in [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 Xiong et al. .

Definition 3 (Policy Improvement Oracle).

For any model $ℳ=(𝒮,𝒜,H,\mathbb{P},d_0,u)$ and a reference function ${\pi }_{\mathrm{ref}}$, we can compute the optimal policy associated with the model ${[{\pi }_{ℳ,h}]}_{h=1}^H$ iteratively as in (7).

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

$${\widehat{u}}_t=\underset{\widehat{u}\in 𝒰}{argmax}{L}_t(\widehat{u}):=\sum _{(x,{\tau }^1,{\tau }^2,z)\in {\cup }_{i=1}^{t-1}{𝒟}_i}\left[z\log (\sigma (\widehat{u}({\tau }^1)-\widehat{u}({\tau }^2)))+(1-z)\log (\sigma (\widehat{u}({\tau }^2)-\widehat{u}({\tau }^1)))\right],$$ (18)

and also an MLE of the transition kernel ${\mathbb{P}}^{\ast }$ over a set $𝒫$ as

$${\widehat{\mathbb{P}}}_t=\underset{\widehat{\mathbb{P}}\in 𝒫}{argmax}{L}_t(\widehat{\mathbb{P}}):=\sum _{(\pi ,\tau )\in {\cup }_{i=1}^{t-1}{𝒟}_i}\log {\widehat{\mathbb{P}}}^{\pi }(\tau ),$$ (19)

where ${\mathbb{P}}^{\pi }(\tau )$ denotes the probability of trajectory $\tau$ under policy $\pi$ and transition kernel $\mathbb{P}$. With the obtained model ${\widehat{ℳ}}_t=({\widehat{u}}_t,{\widehat{\mathbb{P}}}_t)$, the Oracle defined in Definition 3 is called with the reference policy ${\pi }_{\mathrm{ref}}$ set as the initial policy ${\pi }_0$, whose output is adopted as the main policy ${\pi }_t^1$.

Then, we specify how to choose a theoretically sound exploration policy ${\pi }_t^2$. The previous work of Xiong et al. 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.

Proposition 1 (Value Decomposition Lemma for KL-regularized MDP).

If considering a set of $Q$-functions ${[{\widehat{Q}}_h]}_{h=1}^H$ and a reference policy ${\pi }_{\mathrm{ref}}$ with the induced policy $\widehat{\pi }$ as

${\widehat{\pi }}_h(a_h|s_h)\propto {\pi }_{\mathrm{ref},h}(a_h|s_h)\cdot \exp \left({\widehat{Q}}_h(s_h,a_h)/\eta \right),$

and the corresponding set of $V$-functions ${[{\widehat{V}}_h]}_{h=1}^H$ as

${\widehat{V}}_h(s_h)={𝔼}_{a_h\sim {\widehat{\pi }}_h(\cdot |s_h)}\left[{\widehat{Q}}_h(s_h,a_h)\right]-\eta D_{\mathrm{KL}}({\widehat{\pi }}_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h)),{\widehat{V}}_{H+1}(s_{H+1})=0,$

for any comparator policy $\pi$, it holds that

$J(\pi )-J(\widehat{\pi })$ $={𝔼}_{d_0,\pi ,{\mathbb{P}}^{\ast }}[u^{\ast }(s_H,a_H)]-{𝔼}_{d_0,\widehat{\pi },{\mathbb{P}}^{\ast }}[u^{\ast }(s_H,a_H)]$

$+{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,\pi ,{\mathbb{P}}^{\ast }}\left[{\widehat{V}}_{h+1}(s_{h+1})-{\widehat{Q}}_h(s_h,a_h)\right]-{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,\widehat{\pi },{\mathbb{P}}^{\ast }}\left[{\widehat{V}}_{h+1}(s_{h+1})-{\widehat{Q}}_h(s_h,a_h)\right]$

$-\eta \cdot {\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,\pi ,{\mathbb{P}}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\widehat{\pi }}_h(\cdot |s_h))\right],$

where the expectation ${𝔼}_{d_0,\pi ,{\mathbb{P}}^{\ast }}$ is with respect to the prompt and response (i.e., the trajectory) generated following $d_0,{\mathbb{P}}^{\ast }$ and $\pi$.

Based on Proposition 1, the exploration policy ${\pi }_t^2$ is selected as

${\pi }_t^2=\underset{\pi }{argmax}\underset{\tilde{u}\in {\widetilde{𝒰}}_t,\tilde{\mathbb{P}}\in {\widetilde{𝒫}}_t}{\max }$ $\underset{\text{uncertainty measurement of reward estimation}}{\underbrace{{𝔼}_{d_0,\pi ,\tilde{\mathbb{P}}}\left[\tilde{u}(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,\tilde{\mathbb{P}}}\left[\tilde{u}(s_H,a_H)\right]-\left({𝔼}_{d_0,\pi ,\tilde{\mathbb{P}}}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,\tilde{\mathbb{P}}}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)}}$

$+\underset{\text{uncertainty measurement of transition estimation}}{\underbrace{{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,\pi ,\tilde{\mathbb{P}}}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]}},$ (20)

where ${\widetilde{𝒰}}_t$ and ${\widetilde{𝒫}}_t$ are two confidence sets defined as

${\widetilde{𝒰}}_t$ $=\{u\in 𝒰:L_t(u)\ge L_t({\widehat{u}}_t)-c_1\log (|𝒰|T/\delta )\},$ (21)

${\widetilde{𝒫}}_t$ $=\{\mathbb{P}\in 𝒫:L_t(\mathbb{P})\ge L_t({\widehat{\mathbb{P}}}_t)-c_1\log (|𝒫|T/\delta )\}$

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 ${\pi }_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 $\mathbb{P}$), the uncertainty from the estimation of transitions dimimishes, which leads to a similar uncertainty measurement adopted in Xiong et al. .

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

Theorem 1.

Assuming $u^{\ast }\in 𝒰$ and ${\mathbb{P}}^{\ast }\in 𝒫$, with probability at least $1-\delta$, we have that

$\text{Reg}(T)\lesssim$ ${\kappa }^{-1}B\sqrt{d_{𝒰}T\log (|𝒰|T/\delta )}+B^{2}H\xi (d_{𝒫},T,c_2\log (|𝒫|HT/\delta ))$

$-\eta \cdot {\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }^{\ast },{\mathbb{P}}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h^{\ast }(\cdot |s_h),{\pi }_{t,h}^1(\cdot |s_h))\right],$

where $\kappa :=1/(2+\exp (-B)+\exp (B))$, $c_2$ is an absolute constant, $d_{𝒰}$ is the Eluder coefficient defined in Definition 4 while $d_{𝒫}$ and $\xi (\cdot )$ are from the generalized Eluder-type condition defined in Definition 5.

We note that the Eluder coefficient and the generalized Eluder-type condition are standard and well-adopted conditions in the theoretical studies on RL (Zhang, 2023; Zhong et al., 2022; Liu et al., 2023a; Xie et al., 2022; Agarwal et al., 2023) and also RLHF (Zhan et al., 2023; Wang et al., 2023b; Ye et al., 2024). Moreover, for a board class of RL problems (see Zhang (2023); Liu et al. (2023a) for more details), the Eluder coefficient $d_{𝒰}$ is small and the condition is satisfied with $\xi (d_{𝒫},T,c_2\log (|𝒫|HT/\delta ))\lesssim \sqrt{d_{𝒫}T\log (|𝒫|HT/\delta )}$, which implies that the regret of theoretical version of Algorithm 1 is sublinear in $T$, further evidencing its statistical efficiency.

3 Experiments

3.1 Experiment Setup

Task, and datasets.

We use the test sets of MATH (Hendrycks et al., 2021) and GSM8K (Cobbe et al., 2021a) 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 Gou et al. (2023b); Yu et al. (2023); Yue et al. (2023); Liu and Yao (2024); Toshniwal et al. (2024) 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 (Yu et al., 2023) and MMIQC (Liu and Yao, 2024). 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 (Liu and Yao, 2024). 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 (Team et al., 2024), CodeGemma-1.1-it-7B (Team, 2024), Mistral-7B-v0.3 (Jiang et al., 2023), 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 $\backslash$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 1 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 (Bai et al., 2022; Dong et al., 2024). 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 (Gou et al., 2023b), using a subset of the Open-MathInstruct dataset, which was generated by the permissively licensed Mixtral-8x7B 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 (Team et al., 2024) and a learning rate of 1e-5 for Mistral-v0.3 model (Jiang et al., 2023). 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 (Team et al., 2023) or GPT4 (OpenAI, 2023) to further verify the correctness of the trajectory step by step or construct a PRM (Lightman et al., 2023; Wang et al., 2023a) 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 (Lin et al., 2023). 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 ${\lambda }_{+}={\lambda }_{-}=1$ following the original KTO paper (Ethayarajh et al., 2024). 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 B due to space constraint.

3.2 Main Results

We evaluate the models in the zero-shot setting and report the main results in Table 1.

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 Toshniwal et al. (2024) 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 Luo et al. (2023). We also include the reward ranked fine-tuning (RAFT) as a baseline (Dong et al., 2023), which is also known as rejection sampling fine-tuning in the literature (Touvron et al., 2023). 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 (Rafailov et al., 2023; Ethayarajh et al., 2024), 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.

Table 1: Main results of different methods on the test sets of GSM8K and MATH. The SFT training with external tool is based on (a subset of) Open-MathInstruct so the results are generally comparable to the previous SFT models. $†$: the model also serves as the starting checkpoint of other methods except for prompting and CoT without tool use. All the models are allowed to use code interpreter except for the CoT without tool use. The results of the CoT methods are borrowed from the technical reports (Toshniwal et al., 2024; Gou et al., 2023b). The gains relative to the SFT starting checkpoint is marked by  $\uparrow$.

Base Model	Method	with Tool	GSM8K	MATH	AVG
WizardMath-7B	SFT for CoT	✗	54.9	10.7	32.8
WizardMath-13B	SFT for CoT	✗	63.9	14.0	39.0
WizardMath-70B	SFT for CoT	✗	81.6	22.7	52.2
CodeLLaMA-2-7B	SFT	✓	75.9	43.6	59.8
CodeLLaMA-2-13B	SFT	✓	78.8	45.5	62.2
CodeLLaMA-2-34B	SFT	✓	80.7	48.3	64.5
LLaMA-2-70B	SFT	✓	84.7	46.3	65.5
CodeLLaMA-2-70B	SFT	✓	84.6	50.7	67.7
Gemma-1.1-it-7B	SFT†	✓	77.5	46.1	61.8
Gemma-1.1-it-7B	RAFT	✓	79.2	47.3	63.3
Gemma-1.1-it-7B	Iterative Single-turn DPO	✓	81.7	48.9	65.3
Gemma-1.1-it-7B	Iterative Single-turn KTO	✓	80.6	49.0	64.8
Gemma-1.1-it-7B	Iterative M-DPO + fixed reference	✓	79.9	48.0	64.0
Gemma-1.1-it-7B	M-DPO Iteration 1	✓	81.5	49.1	65.3
Gemma-1.1-it-7B	M-DPO Iteration 2	✓	82.5	49.7	66.1
Gemma-1.1-it-7B	M-DPO Iteration 3	✓	83.9  $\uparrow$6.4	51.2  $\uparrow$5.1	67.6  $\uparrow$5.8
Gemma-1.1-it-7B	Iterative M-KTO	✓	82.1  $\uparrow$4.6	49.5  $\uparrow$3.4	65.8  $\uparrow$4.0
CodeGemma-1.1-it-7B	SFT†	✓	77.3	46.4	61.9
CodeGemma-1.1-it-7B	RAFT	✓	78.8	48.4	63.6
CodeGemma-1.1-it-7B	Iterative Single-turn DPO	✓	79.1	48.9	64.0
CodeGemma-1.1-it-7B	Iterative Single-turn KTO	✓	80.2	48.6	64.4
CodeGemma-1.1-it-7B	Iterative M-DPO	✓	81.5  $\uparrow$4.2	50.1  $\uparrow$3.7	65.8  $\uparrow$4.0
CodeGemma-1.1-it-7B	Iterative M-KTO	✓	81.6  $\uparrow$4.3	49.6  $\uparrow$3.2	65.6  $\uparrow$3.8
Mistral-7B-v0.3	SFT†	✓	77.8	42.7	60.3
Mistral-7B-v0.3	RAFT	✓	79.8	43.7	61.8
Mistral-7B-v0.3	Iterative Single-turn DPO	✓	79.8	45.1	62.5
Mistral-7B-v0.3	Iterative Single-turn KTO	✓	81.3	46.3	63.8
Mistral-7B-v0.3	Iterative M-DPO	✓	82.3  $\uparrow$4.5	47.5  $\uparrow$4.8	64.9  $\uparrow$4.7
Mistral-7B-v0.3	Iterative M-KTO	✓	81.7  $\uparrow$3.9	46.7  $\uparrow$4.0	64.2  $\uparrow$4.0
Gemma-2-it-9B	SFT†	✓	84.1	51.0	67.6
Gemma-2-it-9B	RAFT	✓	84.2	52.6	68.4
Gemma-2-it-9B	Iterative Single-turn DPO	✓	85.2	53.1	69.2
Gemma-2-it-9B	Iterative Single-turn KTO	✓	85.4	52.9	69.2
Gemma-2-it-9B	Iterative M-DPO	✓	86.3  $\uparrow$2.2	54.5  $\uparrow$3.5	70.4  $\uparrow$2.9
Gemma-2-it-9B	Iterative M-KTO	✓	86.1  $\uparrow$2.0	54.5  $\uparrow$3.5	70.3  $\uparrow$2.8

From the first two sections in Table 1, 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 (Dong et al., 2023; Yuan et al., 2023a). 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 (Anthony et al., 2017). 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) (Xiong et al., ; Xu et al., 2023) 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 3, 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.

Figure 3: An example of external messages returned by the Python interpreter. The model writes down a bad python code leading to an anomalous and lengthy error message.

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.

Figure 4: The pass@n rate with respect to the number of candidates n. We evaluate the models using temperature 0.7 following the previous works Shao et al. (2024); Toshniwal et al. (2024). We notice that preference learning only improves the metric pass@n when n is relatively small.

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 4. 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 Shao et al. (2024), 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.

3.3 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 Tunstall et al. (2023) 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 ${\pi }_0$. Meanwhile, for online iterative training, Xiong et al. ; Dong et al. (2024) adopt a fixed reference model of ${\pi }_0$, and achieves continuous improvements as the training iterates. In our ablation study, we consider two different choices: (1) using the fixed reference model ${\pi }_0$; (2) updating the reference model to the last iteration’s model at each round. Moreover, we search the KL coefficient $\eta \in \{0.01,0.1,0.5\}$. The results are summarized in Table 2. First, we notice that if we update the reference model at each iteration, the final model outperforms the one with a fixed reference model ${\pi }_0$ with a large margin. Essentially, this dynamic approach optimizes the non-regularized reward, while the one with a fixed reference model ${\pi }_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 5 along the way of iterative training. As we can see, for the first iteration, the results align with Tunstall et al. (2023)’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.

Table 2: Ablation study of the impact of KL regularization. The SFT policy is the starting checkpoint for all other experiments.

Model	Method	GSM8K	MATH
Gemma-1.1-it-7B	SFT 3 epoch	77.5	46.1
Gemma-1.1-it-7B	Iterative M-DPO + $\eta =0.01$	81.7	50.1
Gemma-1.1-it-7B	Iterative M-DPO + $\eta =0.1$	83.9	51.2
Gemma-1.1-it-7B	Iterative M-DPO + $\eta =0.5$	82.8	49.7
Gemma-1.1-it-7B	Iterative M-DPO + fixed reference + $\eta =0.1$	79.9	48.0

Figure 5: The plot of test accuracy on GSM8K dataset and iterations with different levels of KL regularization.

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 (Dong et al., 2024) 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 5, 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 6. 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 (Bai et al., 2022; Touvron et al., 2023; Xiong et al., ; Pace et al., 2024; Dong et al., 2024). 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 4 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\in \{6,12,30\}$ and summarize the results in Table 3. 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.

Table 3: Ablation study of the impact of sampling strategy. The SFT policy is the starting checkpoint for all other experiments. Mixture sampling is adopted for the iterative M-DPO training by default and we run for three iterations in total.

Model	Method	GSM8K	MATH
Gemma-1.1-it-7B	SFT 3 epoch	77.5	46.1
Gemma-1.1-it-7B	Iterative M-DPO with N=30	83.9	51.2
Gemma-1.1-it-7B	Iterative M-DPO with N=12	83.5	51.2
Gemma-1.1-it-7B	Iterative M-DPO with N=6	82.0	49.2
Gemma-1.1-it-7B	Iterative M-DPO with N=30 + On-policy sampling	83.1	49.5

Figure 6: The plot of test accuracy on MATH dataset in terms of training iterations with different sampling strategies.

The best model is obtained with starting checkpoint fine-tuned with more than 1 epochs.

Tunstall et al. (2023) 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 4. 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 (Tunstall et al., 2023) or the iterative DPO for general chatbot (Dong et al., 2024) with only one epoch of SFT.

Table 4: Ablation study of the impact of SFT epoch. Mixture sampling is adopted for the iterative M-DPO training and we run for three iterations in total. The gains relative to their starting SFT checkpoints are marked by  $\uparrow$.

Model	Method	GSM8K	MATH
Gemma-1.1-it-7B	SFT 1 epoch	75.1	41.1
Gemma-1.1-it-7B	SFT 1 epoch + Iterative M-DPO	80.6  $\uparrow$5.5	46.7  $\uparrow$5.6
Gemma-1.1-it-7B	SFT 2 epoch	75.3	44.0
Gemma-1.1-it-7B	SFT 2 epoch + Iterative M-DPO	82.4  $\uparrow$7.1	49.8  $\uparrow$5.8
Gemma-1.1-it-7B	SFT 3 epoch	77.5	46.1
Gemma-1.1-it-7B	SFT 3 epoch + Iterative M-DPO	83.9  $\uparrow$6.4	51.2  $\uparrow$5.1

NLL loss helps when the SFT model is substantially underfitting.

The recent work Pang et al. (2024) 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 (Rafailov et al., 2023) tends to reduce the likelihood of the preferred responses, and this is believed to hurt the reasoning ability (Wang et al., 2024). 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 5. 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 Pang et al. (2024) 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 (Liu et al., 2024b). 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 Pang et al. (2024) 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 Pang et al. (2024), who utilized a general LLaMA2-70B-chat model (Touvron et al., 2023) without firstly fine-tuning on the in-domain data. Our observations align with prior research in the context of developing general chatbots (Lin et al., 2023), which suggests that RLHF is less effective without preliminary SFT.

Table 5: Other ablation studies. Mixture sampling is adopted for the iterative M-DPO training and we run for three iterations in total. The gains relative to the iterative M-DPO are marked by  $\uparrow$.

Model	Method	GSM8K	MATH
Gemma-1.1-it-7B	SFT 3 epoch	77.5	46.1
Gemma-1.1-it-7B	SFT 3 epoch + Iterative M-DPO	83.9	51.2
Gemma-1.1-it-7B	Iterative M-DPO with NLL loss	81.7  $\downarrow$2.2	49.5  $\downarrow$1.7
Gemma-1.1-it-7B	SFT 100 steps	50.8	23.7
Gemma-1.1-it-7B	+ M-DPO Iteration 1	57.8	27.9
Gemma-1.1-it-7B	+ M-DPO and NLL loss Iteration 1	61.0  $\uparrow$3.2	30.1  $\uparrow$2.2

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 (Yuan et al., 2024; Hong et al., 2024; Meng et al., 2024). 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:

${\nabla }_{\theta }{ℒ}_{DPO}({\pi }_{\theta },{\pi }_{\mathrm{ref}})=-\eta \cdot \sigma \left(r_{\theta }(x,y^l)-r_{\theta }(x,y^w)\right)\left[{\displaystyle \frac{1}{{\pi }_{\theta }(y^w|x)}}{\nabla }_{\theta }{\pi }_{\theta }(y^w|x)-{\displaystyle \frac{1}{{\pi }_{\theta }(y^l|x)}}{\nabla }_{\theta }{\pi }_{\theta }(y^l|x)\right],$

where $r_{\theta }(x,y)=\eta \log \frac{{\pi }_{\theta }(x,y)}{{\pi }_{\mathrm{ref}}(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 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 Guo et al. (2024a). 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.

4 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 (Cobbe et al., 2021a) and MATH (Hendrycks et al., 2021). 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 (Lightman et al., 2023) 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 (Pace et al., 2024; Dong et al., 2024; Xu et al., 2023; Hoang Tran, 2024) and Monte Carlo tree search (MCTS) (Xie et al., 2024b). Furthermore, it is also possible to leverage some well-established tricks like adaptive margin and length regularization for DPO training (Hong et al., 2024; Meng et al., 2024). 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 (Munos et al., 2023; Ye et al., 2024). 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.

*

Appendix A Notation Table

Notation	Description
$x,𝒳$	The prompt and the prompt space.
$d_0$	The distribution of initial state (prompt).
$s_h\in 𝒮,a_h\in 𝒜,o_h$	The state, action, and observation.
$H$	Episode length, e.g., the maximal number of tool calls.
${\mathbb{P}}^{\ast }={\left[{\mathbb{P}}_h^{\ast }\right]}_{h=1}^H$	The true observation kernel.
$\tau =(x,y)$	$\tau$ is a trajectory and $y$ is the completion part, i.e., we exclude $x$ from $\tau$.
$u^{\ast }$	The true utility function associated with the BT model defined in Definition 1.
${ℳ}^{\ast }=(𝒮,𝒜,H,{\mathbb{P}}^{\ast },d_0,u^{\ast })$	The true model with observation kernel ${\mathbb{P}}^{\ast }$ and utility function $u^{\ast }$
$\sigma (\cdot )$	$\sigma \left(z\right)=1/\left(1+\exp \left(-z\right)\right)$ is the sigmoid function.
$z\in \{0,1\}$	Preference signal.
$\pi ={\left[{\pi }_h\right]}_{h=1}^H$	The policy, which is parameterized by the LLM.
$ℳ=(𝒮,𝒜,H,\mathbb{P},d_0,u)$	One arbitrary environment with observation kernel $\mathbb{P}$ and utility function $u$.
${\pi }_{\mathrm{ref}}={\left[{\pi }_{\mathrm{ref},h}\right]}_{h=1}^H$	One arbitrary reference policy.
$J(\pi ;ℳ,{\pi }_{\mathrm{ref}})$	The KL-regularized target ((5)) with environment $ℳ$ and reference ${\pi }_{\mathrm{ref}}$.
$\eta$	The coefficient of KL penalty, defined in (5).
$Q_{ℳ}={\left[Q_{ℳ,h}\right]}_{h=1}^H$	The optimal $Q$-values associated with $J(\pi ;ℳ,{\pi }_{\mathrm{ref}})$, defined in (6).
$V_{ℳ}={\left[V_{ℳ,h}\right]}_{h=1}^H$	The optimal $V$-values associated with $J(\pi ;ℳ,{\pi }_{\mathrm{ref}})$, defined in (7).
${\pi }_{ℳ}={\left[{\pi }_{ℳ,h}\right]}_{h=1}^H$	The optimal policy associated with $J(\pi ;ℳ,{\pi }_{\mathrm{ref}})$, defined in (7).
${ℒ}_{\text{M-DPO}}(\cdot )$	M-DPO loss, defined in (12).
${ℒ}_{\text{M-KTO}}(\cdot )$	M-KTO loss, defined in (13).
$J\left(\pi \right)$	The abbreviation of $J(\pi ;{ℳ}^{\ast },{\pi }_0)$, defined in (16).
${\pi }^{\ast }={\left[{\pi }_h^{\ast }\right]}_{h=1}^H$	The optimal policy associated with $J\left(\pi \right)$.
${\pi }_t^1,{\pi }_t^2$	The main and exploration policy at round $t$
$\text{Reg}\left(T\right)$	Regret over horizon $T$, defined in (17).
$𝒰,𝒫$	Known sets such that $u^{\ast }\in 𝒰$ and ${\mathbb{P}}^{\ast }\in 𝒫$
$B$	Assuming $u^{\ast }(x,y)\in [0,B],\forall (x,y)$.
${\widehat{u}}_t,{\widehat{\mathbb{P}}}_t$	MLE of $u^{\ast }$ and ${\mathbb{P}}^{\ast }$ at round $t$, defined in (18) and (19).
${\widetilde{𝒰}}_t,{\widetilde{𝒫}}_t$	Confidences sets of $u^{\ast }$ and ${\mathbb{P}}^{\ast }$ at round $t$, defined in (21).
$c_1,c_2,c$	Absolute constants.
$\kappa$	$1/\left(2+\exp \left(-B\right)+\exp \left(B\right)\right)$.
$d_{𝒰}$	Eluder coefficient from Definition 4.
$d_{𝒫},\xi (\cdot )$	Generalized Eluder-type condition from Definition 5.
$\mathrm{TV}(\cdot ,\cdot )$	Total variation distance between two distributions.

Table 6: The table of notations used in this paper.

Appendix B Implementation Detail

Tools in Math Problem Solving.

Following Gou et al. (2023b); Toshniwal et al. (2024), 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 IPython, 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 Toshniwal et al. (2024). We notice that some works (e.g. Shao et al. (2024)) 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 Gou et al. (2023b); Toshniwal et al. (2024). For evaluating the models with pass@n rate, we follow Toshniwal et al. (2024) 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 (Dong et al., 2024).

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\in \{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 1.

Appendix C Omitted Theoretical Proofs

C.1 Proof of Proposition 1

Proof of Proposition 1.

For one policy $\pi$, starting with $V_{ℳ,H+1}^{\pi }=0$, we recursively define its $V$-value and $Q$-value functions on one model $ℳ=(𝒮,𝒜,H,\mathbb{P},d_0,u)$ and the reference policy ${\pi }_{\mathrm{ref}}$ as

$Q_{ℳ,h}^{\pi }(s_h,a_h)$

$V_{ℳ,h}^{\pi }(s_h)$ $:={𝔼}_{a_h\sim {\pi }_h(\cdot |s_h)}[Q_{ℳ,h}^{\pi }(s_h,a_h)-\eta \cdot D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))].$

It is noted that with the optimal policy ${\pi }_{ℳ}$, $Q_{ℳ,h}=Q_{ℳ,h}^{{\pi }_{ℳ}}$ and $V_{ℳ,h}=V_{ℳ,h}^{{\pi }_{ℳ}}$. In the following discussions, we exclusively focus on the model ${ℳ}^{\ast }=(𝒮,𝒜,H,{\mathbb{P}}^{\ast },d_0,u^{\ast })$ with abbreviations $Q_h^{\pi }=Q_{{ℳ}^{\ast },h}^{\pi }$ and $V_h^{\pi }=V_{{ℳ}^{\ast },h}^{\pi }$.

For any comparator policy $\pi$, it holds that

$J(\pi )-J(\widehat{\pi })={𝔼}_{d_0}\left[V_1^{\pi }(s_1)-{\widehat{V}}_1(s_1)\right]-{𝔼}_{d_0}\left[V_1^{\widehat{\pi }}(s_1)-{\widehat{V}}_1(s_1)\right],$

For any $h\in [H]$, we can obtain that

${𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[V_h^{\pi }(s_h)-{\widehat{V}}_h(s_h)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[V_h^{\widehat{\pi }}(s_h)-{\widehat{V}}_h(s_h)\right]$

$\stackrel{(a)}{=}{𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}[{𝔼}_{{\pi }_h}\left[Q_h^{\pi }(s_h,a_h)\right]-\eta D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))]$

$-{𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}[{𝔼}_{{\widehat{\pi }}_h}\left[{\widehat{Q}}_h(s_h,a_h)\right]-\eta D_{\mathrm{KL}}({\widehat{\pi }}_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))]$

$-{𝔼}_{d_0,{\widehat{\pi }}_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}[{𝔼}_{{\widehat{\pi }}_h}\left[Q_h^{\widehat{\pi }}(s_h,a_h)\right]-\eta D_{\mathrm{KL}}({\widehat{\pi }}_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))]$

$+{𝔼}_{d_0,{\widehat{\pi }}_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}[{𝔼}_{{\widehat{\pi }}_h}\left[{\widehat{Q}}_h(s_h,a_h)\right]-\eta D_{\mathrm{KL}}({\widehat{\pi }}_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))]$

$={𝔼}_{d_0,{\pi }_{1:h},{\mathbb{P}}_{1:h-1}^{\ast }}\left[Q_h^{\pi }(s_h,a_h)-{\widehat{Q}}_h(s_h,a_h)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:h},{\mathbb{P}}_{1:h-1}^{\ast }}\left[Q_h^{\widehat{\pi }}(s_h,a_h)-{\widehat{Q}}_h(s_h,a_h)\right]$

$+{𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\underset{\text{term (I)}}{\underbrace{\left[{𝔼}_{{\pi }_h}\left[{\widehat{Q}}_h(s_h,a_h)\right]-{𝔼}_{{\widehat{\pi }}_h}\left[{\widehat{Q}}_h(s_h,a_h)\right]\right]}}$

$-\eta \cdot {𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))\right]+\eta \cdot {𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[D_{\mathrm{KL}}({\widehat{\pi }}_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))\right]$

$\stackrel{(b)}{=}{𝔼}_{d_0,{\pi }_{1:h},{\mathbb{P}}_{1:h-1}^{\ast }}\left[Q_h^{\pi }(s_h,a_h)-{\widehat{Q}}_h(s_h,a_h)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:h},{\mathbb{P}}_{1:h-1}^{\ast }}\left[Q_h^{\widehat{\pi }}(s_h,a_h)-{\widehat{Q}}_h(s_h,a_h)\right]$

$-\eta \cdot {𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\widehat{\pi }}_h(\cdot |s_h))\right].$

In the above derivation, equation (a) is from the definitions of $Q^{\pi }$ and $V^{\pi }$, and the relationship between $\widehat{Q}$ and $\widehat{V}$. The equation (b) is because

(term I) $:={𝔼}_{{\pi }_h}\left[{\widehat{Q}}_h(s_h,a_h)\right]-{𝔼}_{{\widehat{\pi }}_h}\left[{\widehat{Q}}_h(s_h,a_h)\right]$

$=\eta \cdot {𝔼}_{{\pi }_h}\left[\log {\displaystyle \frac{{\widehat{\pi }}_h(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}\right]-\eta \cdot {𝔼}_{{\widehat{\pi }}_h}\left[\log {\displaystyle \frac{{\widehat{\pi }}_h(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}\right]$

$=\eta \cdot D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h))-\eta \cdot D_{\mathrm{KL}}({\pi }_h(\cdot |s_h),{\widehat{\pi }}_h(\cdot |s_h))-\eta \cdot D_{\mathrm{KL}}({\widehat{\pi }}_h(\cdot |s_h),{\pi }_{\mathrm{ref},h}(\cdot |s_h)).$

where the second equation is from the relationship that

${\widehat{Q}}_h(s_h,a_h)=\eta \cdot \log {\displaystyle \frac{{\widehat{\pi }}_h(a_h|s_h)}{{\pi }_{\mathrm{ref},h}(a_h|s_h)}}-\eta \cdot \log {\widehat{Z}}_h(s_h).$

Furthermore, if $h=H$, we can obtain that

${𝔼}_{d_0,{\pi }_{1:H-1},{\mathbb{P}}_{1:H-1}^{\ast }}\left[V_H^{\pi }(s_H)-{\widehat{V}}_H(s_H)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:H-1},{\mathbb{P}}_{1:H-1}^{\ast }}\left[V_H^{\widehat{\pi }}(s_H)-{\widehat{V}}_H(s_H)\right]$

$={𝔼}_{d_0,{\pi }_{1:H},{\mathbb{P}}_{1:H-1}^{\ast }}\left[u^{\ast }(s_H,a_H)-{\widehat{Q}}_H(s_H,a_H)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:H},{\mathbb{P}}_{1:H-1}^{\ast }}\left[u^{\ast }(s_H,a_H)-{\widehat{Q}}_H(s_H,a_H)\right]$

$-\eta \cdot {𝔼}_{d_0,{\pi }_{1:H-1},{\mathbb{P}}_{1:H-1}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_H(\cdot |s_H),{\widehat{\pi }}_H(\cdot |s_H))\right]$

$={𝔼}_{d_0,{\pi }_{1:H},{\mathbb{P}}_{1:H-1}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:H},{\mathbb{P}}_{1:H-1}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]$

$+{𝔼}_{d_0,{\pi }_{1:H},{\mathbb{P}}_{1:H}^{\ast }}\left[{\widehat{V}}_{H+1}(s_{H+1})-{\widehat{Q}}_H(s_H,a_H)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:H},{\mathbb{P}}_{1:H}^{\ast }}\left[{\widehat{V}}_{H+1}(s_{H+1})-{\widehat{Q}}_H(s_H,a_H)\right]$

$-\eta \cdot {𝔼}_{d_0,{\pi }_{1:H-1},{\mathbb{P}}_{1:H-1}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_H(\cdot |s_H)||{\widehat{\pi }}_H(\cdot |s_H))\right],$

where the second equality leverages that ${\widehat{V}}_{H+1}(s_{H+1})=0$; otherwise, for all $h\le H-1$, it holds that

${𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[V_h^{\pi }(s_h)-{\widehat{V}}_h(s_h)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[V_h^{\widehat{\pi }}(s_h)-{\widehat{V}}_h(s_h)\right]$

$={𝔼}_{d_0,{\pi }_{1:h},{\mathbb{P}}_{1:h-1}^{\ast }}\left[Q_h^{\pi }(s_h,a_h)-{\widehat{Q}}_h(s_h,a_h)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:h},{\mathbb{P}}_{1:h-1}^{\ast }}\left[Q_h^{\widehat{\pi }}(s_h,a_h)-{\widehat{Q}}_h(s_h,a_h)\right]$

$-\eta \cdot {𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h(\cdot |s_h)||{\widehat{\pi }}_h(\cdot |s_h))\right]$

$={𝔼}_{d_0,{\pi }_{1:h},{\mathbb{P}}_{1:h}^{\ast }}\left[{\widehat{V}}_{h+1}(s_{h+1})-{\widehat{Q}}_h(s_h,a_h)\right]-{𝔼}_{d_0,{\widehat{\pi }}_{1:h},{\mathbb{P}}_{1:h}^{\ast }}\left[{\widehat{V}}_{h+1}(s_{h+1})-{\widehat{Q}}_h(s_h,a_h)\right]$

$-\eta \cdot {𝔼}_{d_0,{\pi }_{1:h-1},{\mathbb{P}}_{1:h-1}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h(\cdot |s_h)||{\widehat{\pi }}_h(\cdot |s_h))\right]$

$+{𝔼}_{d_0,{\pi }_{1:h},{\mathbb{P}}_{1:h}^{\ast }}\left[V_{h+1}^{\pi }(s_{h+1})-{\widehat{V}}_{h+1}(s_{h+1})\right]-{𝔼}_{d_0,{\pi }_{1:h},{\mathbb{P}}_{1:h}^{\ast }}\left[V_{h+1}^{\widehat{\pi }}(s_{h+1})-{\widehat{V}}_{h+1}(s_{h+1})\right].$

The proposition can be obtained by iteratively using the above relationship for $h\in [H]$. ∎

C.2 Proof of Theorem 1

First, with the assumption $u^{\ast }\in 𝒰$ and ${\mathbb{P}}^{\ast }\in 𝒫$, the following lemma demonstrates that ${\widetilde{𝒰}}_t$ and ${\widetilde{𝒫}}_t$ are valid confidence sets.

Lemma 1 (Proposition B.1 from Liu et al. (2023a)).

There exists an absolute constant $c_1$ such that for any $\delta \in (0,1]$, with probability at least $1-\delta$, for all $t\in [T]$, $\widehat{u}\in 𝒰$, and $\widehat{\mathbb{P}}\in 𝒫$, it holds that

$L_t(\widehat{u})-L_t(u^{\ast })\le c_1\log (|𝒰|T/\delta ),L_t(\widehat{\mathbb{P}})-L_t({\mathbb{P}}^{\ast })\le c_1\log (|𝒫|T/\delta ),$

which implies that $u^{\ast }\in {\widetilde{𝒰}}_t$ and ${\mathbb{P}}^{\ast }\in {\widetilde{𝒫}}_t$.

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

Lemma 2.

There exists an absolute constant $c_2$ such that for any $\delta \in (0,1]$, with probability at least $1-\delta$, for all $t\in [T]$, we have

and for all $t\in [T]$, $h\in [H]$, we have

where $\mathrm{TV}(\{d_0,\pi ,\mathbb{P}\},\{d_0,{\pi }^{\prime },{\mathbb{P}}^{\prime }\})$ denotes the TV distance between the probability distributions over the trajectories induced by $d_0,\pi ,\mathbb{P}$ and $d_0,{\pi }^{\prime },{\mathbb{P}}^{\prime }$.

Proof of Lemma 2.

First, for ${\tilde{u}}_t$, we can obtain that with probability at least $1-\delta$, there exists an absolute constant $c$ such that for all $t\in [T]$,

$=c\left(L_t(u^{\ast })-L_t({\tilde{u}}_t)+\log (|𝒰|T/\delta )\right)$

$\le c\left(L_t(u^{\ast })-L_t({\widehat{u}}_t)+c_1\log (|𝒰|T/\delta )+\log (|𝒰|T/\delta )\right)$

$\le c_2\log (|𝒰|T/\delta ).$

where the first inequality is from Proposition B.2 from Liu et al. (2023a) and the second inequality uses Lemma 1. The result for ${\widehat{u}}_t$ can be similarly established.

Then, following similar steps, for ${\tilde{\mathbb{P}}}_t$, we can obtain that with probability at least $1-\delta$, there exists an absolute constant $c$ such that for all $t\in [T]$,

$=c\cdot \left(L_t({\mathbb{P}}^{\ast })-L_t({\tilde{\mathbb{P}}}_t)+2\log (|𝒫|HT/\delta )\right)$

$\le c\cdot \left(L_t({\mathbb{P}}^{\ast })-L_t({\widehat{\mathbb{P}}}_t)+c_1\log (|𝒫|T/\delta )+2\log (|𝒫|HT/\delta )\right)$

$\le c_2\log (|𝒫|HT/\delta ).$

The result for ${\widehat{\mathbb{P}}}_t$ can also be similarly established. ∎

Proof of Theorem 1.

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

$J({\pi }^{\ast })-J({\pi }_t^1)$

$={𝔼}_{d_0,{\pi }^{\ast },{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }^{\ast },{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)$

$+{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }^{\ast },{\mathbb{P}}^{\ast }}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]-{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]$

$\le \underset{{\text{term (I)}}_t}{\underbrace{{𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)}}$

$+\underset{{\text{term (II)}}_t}{\underbrace{{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]+{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)-{\widehat{V}}_{t,h+1}(s_{h+1})\right]}},$

where the inequality is from the definition of ${\pi }_t^2$ and the fact that $(u^{\ast },{\mathbb{P}}^{\ast })\in {\widetilde{𝒰}}_t\times {\widetilde{𝒫}}_t$ from Lemma 1.

We define the following terms:

${\text{term (A)}}_t$ $:={𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]\right),$

${\text{term (B)}}_t$ $:={𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]\right),$

${\text{term (C)}}_t$ $:={\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^j,{\mathbb{P}}^{\ast }}\left[\mathrm{TV}({\tilde{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h))\right],$

${\text{term (D)}}_t$ $:={\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^j,{\mathbb{P}}^{\ast }}\left[\mathrm{TV}({\widehat{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h))\right].$

For ${\text{term (I)}}_t$, we have that

${\text{term (I)}}_t$ $:={𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)$

$={𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]\right)$

$+{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)$

$+{𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]\right)$

$+{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)$

$\le {𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]\right)$

$+{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)$

$+4B\cdot \mathrm{TV}(\{d_0,{\pi }_t^1,{\tilde{\mathbb{P}}}_t\},\{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }\})+4B\cdot \mathrm{TV}(\{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t\},\{d_0,{\pi }_t^2,\mathbb{P}\})$

$\le \underset{{\text{term (A)}}_t}{\underbrace{{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]\right)}}$

$+\underset{{\text{term (B)}}_t}{\underbrace{{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u_t^{\ast }(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\widehat{u}}_t(s_H,a_H)\right]\right)}}$

$+4B\cdot \underset{{\text{term (C)}}_t}{\underbrace{{\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0}{𝔼}_{{\pi }_t^j,{\mathbb{P}}^{\ast }}\left[\mathrm{TV}({\tilde{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h))\right]}}.$

For ${\text{term (II)}}_t$, we have that

${\text{term (II)}}_t$ $={\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]$

$+{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)-{\widehat{V}}_{t,h+1}(s_{h+1})\right]$

$={\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]$

$+{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]$

$-{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\widehat{V}}_{t,h+1}(s_{h+1})-\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)\right]$

$+{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[\left[{\widehat{\mathbb{P}}}_{t,h}{\widehat{V}}_{t,h+1}\right](s_h,a_h)-{\widehat{V}}_{t,h+1}(s_{h+1})\right]$

$\le 2B\cdot {\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^j,{\mathbb{P}}^{\ast }}[\mathrm{TV}({\widehat{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h)),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)]$

$+2BH\cdot \mathrm{TV}(\{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t\},\{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }\})$

$\le 2B\cdot \underset{{\text{term (D)}}_t}{\underbrace{{\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^j,{\mathbb{P}}^{\ast }}[\mathrm{TV}({\widehat{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h)),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)]}}$

$+2BH\cdot \underset{{\text{term (C)}}_t}{\underbrace{{\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^j,{\mathbb{P}}^{\ast }}[\mathrm{TV}({\tilde{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h)),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)]}}.$

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

$\mathrm{TV}(\{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t\},\{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }\})\le {\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[\mathrm{TV}({\tilde{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h))\right],$

which can be derived as

$\mathrm{TV}(\{d_0,{\pi }_t^2,{\tilde{\mathbb{P}}}_t\},\{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }\})$ $\le {\displaystyle \sum _{h\in [H]}}\mathrm{TV}(\{d_0,{\pi }_t^2,{\mathbb{P}}_{1:h-1}^{\ast },{\tilde{\mathbb{P}}}_{t,h:H}\},\{d_0,{\pi }_t^2,{\mathbb{P}}_{1:h}^{\ast },{\tilde{\mathbb{P}}}_{t,h+1:H}\})$

$={\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[\mathrm{TV}({\tilde{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h)\}\right)].$

Then, we can obtain that

${\displaystyle \sum _{t\in [T]}}J({\pi }^{\ast })-J({\widehat{\pi }}_t^1)$ $\le {\displaystyle \sum _{t\in [T]}}{\text{term (A)}}_t+{\displaystyle \sum _{t\in [T]}}{\text{term (B)}}_t+(4B+2BH){\displaystyle \sum _{t\in [T]}}{\text{term (C)}}_t+2B{\displaystyle \sum _{t\in [T]}}{\text{term (D)}}_t.$

Then, we control the sum of each individual term in the following. First, for ${\text{term (A)}}_t$, with probability at least $1-\delta$, we have that

${\displaystyle \sum _{t\in [T]}}{\text{term (A)}}_t$

$={\displaystyle \sum _{t\in [T]}}{𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[{\tilde{u}}_t(s_H,a_H)\right]-\left({𝔼}_{d_0,{\pi }_t^2,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]-{𝔼}_{d_0,{\pi }_t^1,{\mathbb{P}}^{\ast }}\left[u^{\ast }(s_H,a_H)\right]\right)$

$\le {\displaystyle \sum _{t\in [T]}}{\tilde{u}}_t(s_{t,H}^2,a_{t,H}^2)-{\tilde{u}}_t(s_{t,H}^1,a_{t,H}^1)-\left(u^{\ast }(s_{t,H}^2,a_{t,H}^2)-u^{\ast }(s_{t,H}^1,a_{t,H}^1)\right)+O(B\sqrt{T\log (1/\delta )})$

$\le \sqrt{d_{𝒰}{\displaystyle \sum _{t=2}^T}\left(1+{\displaystyle \sum _{i=1}^{t-1}}{\left({\tilde{u}}_t(s_{i,H}^2,a_{i,H}^2)-{\tilde{u}}_t(s_{i,H}^1,a_{i,H}^1)-\left(u^{\ast }(s_{i,H}^2,a_{i,H}^2)-u^{\ast }(s_{i,H}^1,a_{i,H}^1)\right)\right)}^2\right)}+O(B\sqrt{T\log (1/\delta )})$

$\le \sqrt{d_{𝒰}{\displaystyle \sum _{t=2}^T}\left(1+{\kappa }^{-2}{\displaystyle \sum _{i=1}^{t-1}}{\left(\sigma \left({\tilde{u}}_t(s_{i,H}^2,a_{i,H}^2)-{\tilde{u}}_t(s_{i,H}^1,a_{i,H}^1)\right)-\sigma \left(u^{\ast }(s_{i,H}^2,a_{i,H}^2)-u^{\ast }(s_{i,H}^1,a_{i,H}^1)\right)\right)}^2\right)}+O(B\sqrt{T\log (1/\delta )})$

$\lesssim {\kappa }^{-1}B\sqrt{d_{𝒰}T\log (|𝒰|T/\delta )},$

where the first inequality is from the Hoeffding inequality, the second inequality uses the Eluder coefficient $d_{𝒰}:=\text{EC}(1,𝒰-𝒰,T)$ from Definition 4, the third inequality leverages the mean value theorem with $\kappa :=1/(2+\exp (-B)+\exp (B))$ representing the minimum derivative of $\sigma (\cdot )$ in the regime of $[0,B]$, and the last inequality incorporates Lemma 2. A similar result can be obtained for ${\text{term (B)}}_t$.

For ${\text{term (C)}}_t$, we have that

${\displaystyle \sum _{t\in [T]}}{\text{term (C)}}_t$ $={\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{t\in [T]}}{\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }_t^j,{\mathbb{P}}^{\ast }}\left[\mathrm{TV}({\tilde{\mathbb{P}}}_{t,h}(\cdot |s_h,a_h),{\mathbb{P}}_h^{\ast }(\cdot |s_h,a_h))\right]$

$={\displaystyle \sum _{j\in \{1,2\}}}{\displaystyle \sum _{t\in [T]}}{\displaystyle \sum _{h\in [H]}}\mathrm{TV}(\{d_0,{\pi }_t^j,[{\mathbb{P}}_{1:h-1}^{\ast },{\tilde{\mathbb{P}}}_{t,h},{\mathbb{P}}_{h+1:H}^{\ast }]\},\{d_0,{\pi }_t^j,{\mathbb{P}}_{1:H}^{\ast }\})$

$\le 2H\cdot \xi (d_{𝒫},T,c_2\log (|𝒫|HT/\delta )),$

where the last step is from the generalized Eluder-type condition in Definition 5 and Lemma 2. A similar result can be obtained for ${\text{term (D)}}_t$.

Finally, we obtain that

$\text{Reg}(T)\lesssim$ ${\kappa }^{-1}B\sqrt{d_{𝒰}T\log (|𝒰|T/\delta )}+B^{2}H\xi (d_{𝒫},T,c_2\log (|𝒫|HT/\delta )$

$-\eta \cdot {\displaystyle \sum _{h\in [H]}}{𝔼}_{d_0,{\pi }^{\ast },{\mathbb{P}}^{\ast }}\left[D_{\mathrm{KL}}({\pi }_h^{\ast }(\cdot |s_h),{\pi }_{t,h}^1(\cdot |s_h))\right],$

which concludes the proof. ∎

Appendix D Technical Lemmas

Lemma 3 (Solution of KL-regularized Optimization (Proposition 7.16 and Theorem 15.3 of Zhang (2023))).

Given a loss functional with respect to $p(\cdot |x)$, written as

$${𝔼}_{w\sim p(\cdot )}\left[-U(w)+\eta D_{\mathrm{KL}}(p(\cdot ),p_0(\cdot ))\right]=\eta D_{\mathrm{KL}}(p(\cdot ),p_0(\cdot )\exp \left(\frac{1}{\eta }U(\cdot )\right))-\eta \cdot \log \underset{C_r}{\underbrace{{𝔼}_{w\sim p_0(\cdot )}\exp \left(\frac{1}{\eta }U(w)\right)}},$$

where the minimizer of the loss functional is $p^{\ast }(w)=\frac{1}{C_r}{p}_0(w)\exp \left(\frac{1}{\eta }U(w)\right)$, also known as Gibbs distribution.

Definition 4 (Eluder Coefficient, Definition 17.17 in Zhang (2023)).

Given a function class $ℱ$, its Eluder coefficient $\text{EC}(\lambda ,ℱ,T)$ is defined as the smallest number $d$ so that for any sequence $\{x_t:t\in [T]\}$ and $\{f_t:t\in [T]\}\in ℱ$,

${\displaystyle \sum _{t=2}^T}|f_t(x_t)-f^{\ast }(x_t)|\le \sqrt{d{\displaystyle \sum _{t=2}^T}\left(\lambda +{\displaystyle \sum _{i=1}^{t-1}}{(f_t(x_i)-f^{\ast }(x_i))}^2\right)}.$

Definition 5 (Generalized Eluder-type Condition, Condition 3.1 in Liu et al. (2023a)).

There exists a real number $d_{𝒫}\in {ℝ}^{+}$ and a function $\xi$ such that for any $(T,\mathrm{\Delta })\in \mathbb{N}\times {ℝ}^{+}$, transitions $\{{\mathbb{P}}_t^{\prime }:t\in [T]\}$ and policies $\{{\pi }_t:t\in [T]\}$, we have