Building Math Agents with Multi-Turn Iterative Preference Learning
Under review as a conference paper at ICLR 2025

BUILDING MATH AGENTS WITH MULTI-TURN ITERA-
TIVE PREFERENCE LEARNING

Anonymous authors
Paper under double-blind review

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) reason-
ing. While current methods focus on synthetic data generation and Supervised
Fine-Tuning (SFT), this paper studies the complementary direct preference learn-
ing 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 ex-
ternal tool integration required for tool-integrated mathematical reasoning tasks.
To fill in this gap, we introduce a multi-turn direct preference learning frame-
work, 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 frame-
work is validated through training of various language models using an augmented
prompt set from the GSM8K and MATH datasets. Our results demonstrate sub-
stantial improvements: a supervised fine-tuned Gemma-1.1-it-7B model’s perfor-
mance 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 lan-
guage tasks. Notable models include ChatGPT (OpenAI, 2023), Claude (Anthropic, 2023), and
Gemini (Gemini et al., 2023). However, despite these advances, even the most advanced closed-
source LLMs still struggle with complex reasoning tasks that require multi-turn 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; 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 ca-
pabilities (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 math-
ematical problem solving, existing works primarily focus on synthetic data generation (by strong
teacher models) and supervised fine-tuning (SFT), as seen in ToRA (Gou et al., 2023b), Meta-
MathQA (Yu et al., 2023), MAmmoTH (Yue et al., 2023; 2024), and Open-MathInstruct (Toshniwal
et al., 2024). These 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
standard in the LLM training pipeline (Ouyang et al., 2022; Gemini et al., 2023). Broadly speaking,

1

000
001

002
003
004
005
006
007

008
009
010
011
012
013

014
015
016
017
018

019
020
021
022
023
024

025
026
027
028
029
030

031
032
033
034
035
036

037
038
039
040
041

042
043
044
045
046
047

048
049
050
051
052
053

Under review as a conference paper at ICLR 2025

the RLHF learning paradigm, which was originally designed for aligning LLMs with human values
and preferences, is distinct from SFT as it learns from relative feedback. 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).

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., 2024; Rafailov et al., 2023). Then, we de-
rive the optimality condition of the planning with such an MDP and our findings indicate that when
the external randomness is low, we can develop multi-turn direct alignment algorithms (M-DPO and
M-KTO), where the primary modification is to mask out irrelevant tokens during training. Further-
more, we extend our approach to its online iterative variants, which recent works demonstrated to be
promising (Xiong et al., 2024; 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 will make our models, datasets, and code publicly available for further research and develop-
ment.

2 ALGORITHMS DEVELOPMENT

2.1 PROBLEM FORMULATION

We first formally formulate the tool-integrated reasoning task. At the first step, a prompt x ∈ X is
sampled from some distribution d0 as the initial state s1 = x. Then, at each step h ∈ [H],

• Action: the agent observes the current state sh, which is the history of the first h − 1
interactions with the external environment, and takes an action ah according to some policy
πh(·|sh) ∈ ∆(A).

• Observation: in response to the agent’s action, the environment then returns an observation

oh ∼ P∗

h(·|sh, ah)1 based on the history sh and current action ah.

Then, we transit to a new state, which is the history up to the step h + 1: sh+1 = (sh, ah, oh) =
(x, a1, o1, · · · , ah, oh), and a new step begins. This process repeats for H rounds in total and even-
tually, we collect a trajectory: τ = (x, a1, o1, · · · , oH−1, aH ). We present an example of multi-turn
tool-integrated reasoning in Figure 1. Typically, the action is in the ReAct manner, which consist
of a reasoning step fh and an execution step eh (e.g., writing python code) (Yao et al., 2022). We
mention in passing that such an 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.

To connect the problem with RLHF that learns from relative feedback, we follow Ouyang et al.
(2022); Bai et al. (2022) to assume that we can query the Bradley-Terry model for preference signal.
Definition 1 (Bradley-Terry model). We denote τ /x = y, where the prompt is excluded from the tra-
jectory. We assume that there exists a utility function of the trajectory u∗ such that given (x, y1, y2),
one response y1 is preferred over another response y2, denoted as y1 ≻ y2, with probability

Prob(cid:0)y1 ≻ y2 | x, y1, y2(cid:1) = σ(cid:0)u∗(x, y1) − u∗(x, y2)(cid:1),

(1)

1When there is no ambiguity, the abbreviation sh+1 ∼ P∗

h(·|sh, ah) is also adopted.

2

054

055
056
057
058
059

060
061
062
063
064
065

066
067
068
069
070
071

072
073
074
075
076
077

078
079
080
081
082

083
084
085
086
087
088

089
090
091
092
093
094

095
096
097
098
099
100

101
102
103
104
105

106
107

Under review as a conference paper at ICLR 2025

108
109

110
111
112
113
114
115

116
117
118
119
120
121

122
123
124
125
126

127
128
129
130
131
132

133
134
135
136
137
138

139
140
141
142
143
144

145
146
147
148
149

150
151
152
153
154
155

156
157
158
159
160
161

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 fh and an execution step eh.

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

Here we only assume access to the trajectory-level preference, but not an action-level one. However,
we remark that the utility function itself can be defined in a step-wise manner. Examples of the utility
function include the binary reward from checking final result, outcome-supervised reward models
(Cobbe et al., 2021b), and process-supervised reward model (Lightman et al., 2023).

2.2 PLANNING WITH A MODEL: OPTIMALITY CONDITION AND PRACTICAL ALGORITHM

We develop the main algorithms in this section with the general MDP formulation. Following
Rafailov et al. (2023), we first establish the connection between a model M = (S, A, H, P, d0, u)

3

Under review as a conference paper at ICLR 2025

and its associated optimal policy. In particular, we are interested in the following KL-regularized
planning problem with respect to a reference policy πref :

arg max
π

J(π; M, πref ) = Ex∼d0,ah∼πh(·|sh),oh∼Ph(·|sh,ah)

(cid:104)

u(x, y) − η

H
(cid:88)

(cid:0)πh(·|sh), πref,h(·|sh)(cid:1)(cid:105)
.

DKL

h=1

(2)
In the single-turn case with H = 1, the optimal solution with respect to a utility function u is the
Gibbs distribution (see Lemma 3). Moving toward multi-turn case, we first consider H = 2 to
illustrate the idea. The idea is to take a backward iteration from h = H = 2 to h = 1. Specifically,
when we fix s2 and consider only the step 2, it reduces to the single-turn case:

πM,2(·|s2) = arg max

π2

Ea2∼π2(·|s2)

(cid:16)

u(s2, a2) − η · DKL

(cid:0)π2(·|s2), πref,2(·|s2)(cid:1)(cid:17)

∝ πref,2(·|s2) · exp

(cid:16) u(s2, ·)
η

(cid:17)
.

Then, we can define the value function associated with πM,2 as
(cid:2)u(s2, a2) − ηDKL

VM,2(s2) := Ea2∼πM,2(·|s2)

(cid:0)πM,2(·|s2), πref,2(·|s2)(cid:1)(cid:3)

QM,1(s1, a1) := Eo1∼P1(·|s1,a1) [VM,2(s2)] .

For step 1, since we have determined πM,2, with the definition of QM,1(s1, a1), we have

πM,1(·|s1) = arg max

π1

Ea1∼π1(·|x)

(cid:104)

QM,1(s1, a1) − ηDKL

(cid:0)π1(·|s1), πref,1(·|s1)(cid:1)(cid:105)

∝ πref,1(·|s1) · exp

(cid:16) QM,1(s1, ·)
η

(cid:17)
.

By construction, {πM,h}2
h=1 is optimal as it maximizes the KL-regularized target. For general
MDP, we can repeat the process for H times starting with VM,H+1 = 0 where we recursively define

(cid:40)

QM,h(sh, ah) =

u(sH , aH ),
Eoh∼Ph(·|sh,ah)[VM,h+1(sh+1)],

if h = H,

if h ≤ H − 1,

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

πM,h(ah|sh) :=

1
Zh(sh)

πref,h(ah|sh) · exp

(cid:16) QM,h(sh, ah)
η

(cid:17)

(Gibbs distribution of QM,h)

VM,h(sh) := Eah∼πM,h(·|sh)

(cid:2)QM,h(sh, ah) − η · DKL
(cid:16) QM,h(sh, a′
(cid:17)

h)

(cid:0)πM,h(·|sh), πref,h(·|sh)(cid:1)(cid:3),

= η log Eπref,h(a′

h|sh) exp

η

= η log Zh(sh),

(3)

(4)

ah∈A πref,h(ah|sh) · exp (cid:0) QM,h(sh,ah)

where Zh(sh) = (cid:80)
second equality in the definition of the V -value is from Lemma 3. Then, by definition, [πM,h]H
h=1
is optimal. Essentially, we solve H Gibbs distributions in terms of the Q-values. We remark that the
results are essentially from the entropy-regularized MDPs (Williams & Peng, 1991; Ziebart, 2010).

(cid:1) is the normalization constant. The

η

Multi-turn DPO. According to equation 4, we can solve the Q-values as

QM,h(sh, ah) = η · log

πM,h(ah|sh)
πref,h(ah|sh)

+ VM,h(sh).

(5)

Furthermore, combining equation 5 with the definition of Q-values QM,h in equation 3, we have

Eoh∼Ph(·|sh,ah)VM,h+1(sh+1) = η · log

u(sH , aH ) = η · log

πM,h(ah|sh)
πref,h(ah|sh)
πM,H (aH |sH )
πref,H (aH |sH )

+ VM,h(sh),

if h ≤ H − 1

(6)

+ VM,H (sH ).

Summing over h ∈ [H], we have the following re-parameterization result:

u(sH , aH ) = η

(cid:124)

H
(cid:88)

h=1

log

πM,h(ah|sh)
πref,h(ah|sh)
(cid:123)(cid:122)
(cid:125)
term (A)

+ VM,1(s1)
(cid:123)(cid:122)
(cid:125)
term (B)

(cid:124)

+

(cid:104)

(cid:105)
VM,h+1(sh+1) − Eoh∼Ph(·|sh,ah)VM,h+1(sh+1)

.

(cid:123)(cid:122)
term (C)

(cid:125)

(7)

H−1
(cid:88)

h=1
(cid:124)

4

162
163

164
165
166
167
168
169

170
171
172
173
174
175

176
177
178
179
180

181
182
183
184
185
186

187
188
189
190
191
192

193
194
195
196
197
198

199
200
201
202
203

204
205
206
207
208
209

210
211
212
213
214
215

Under review as a conference paper at ICLR 2025

216
217

218
219
220
221
222
223

224
225
226
227
228
229

230
231
232
233
234

235
236
237
238
239
240

241
242
243
244
245
246

247
248
249
250
251
252

253
254
255
256
257

258
259
260
261
262
263

264
265
266
267
268
269

Here, term (A) is similar to the single-turn case and term (B) will be cancelled for the reward
difference of two samples with the same prompt s1. However, in practice, term (C) is typically not
feasible to directly compute as term (C) is related to the randomness of the external environment.

For the focus of this work, i.e., the tool-integrated mathematical reasoning, luckily the code execu-
tion result is determined by the history (the codes written by the LLMs). This leads to term (C) = 0.
Therefore, we can plug equation 7 into the maximum likelihood estimation of the utility function
with a dataset D consisting of (x, τ w, τ l), to get the following multi-turn DPO (M-DPO) loss:

LM-DPO(θ) = −

(cid:88)

(x,τ w ,τ l)∈D

log σ

(cid:16)

η

H
(cid:88)

(cid:104)

log

h=1

πθ,h(aw
πref,h(aw

h |sw
h )
h |sw
h )

− log

πθ,h(al
πref,h(al

h|sl
h)
h|sl
h)

(cid:105)(cid:17)

,

(8)

Similarly, we can implement M-KTO under deterministic transition. We refer interested readers to
Appendix A for the loss function details.

2.3 ONLINE ITERATIVE TRAINING

We now combine the planning algorithm M-DPO with the online iterative learning framework, as
inspired by its great success in the single-turn case (Xiong et al., 2024; Guo et al., 2024b).

Learning objective. For a more comprehensive understanding of its statistical behavior, we will
consider two different learning objectives. The first objective is a KL-regularized one:

max
π

Ex∼d0

Eah∼π(·|sh),oh∼P∗

h(·|sh,ah)

(cid:104)
u∗(x, y) − η

H
(cid:88)

h=1

(cid:0)π(·|sh), π0(·|sh)(cid:1)(cid:105)

,

DKL

(9)

i.e., maxπ J(π; M∗, π0) where M∗ = (S, A, H, P∗, d0, u∗) is the groundtruth environment and π0
is the initial SFT policy. This target is widely adopted in RLHF and requires us to search for the
optimal policy only at a fixed KL ball centered at the SFT policy π0. In contrast, the second one is
the non-regularized target, i.e., directly optimizing the reward:

max
π

Ex∼d0

Eah∼π(·|sh),oh∼P∗

h(·|sh,ah)

(cid:2)u∗(x, y)(cid:3).

(10)

This target is the standard one in canonical RL studies (Sutton & 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.

Algorithmic framework. We present a general online iterative algorithmic framework in Algo-
rithm 1. This framework is termed as Online Iterative Multi-turn Gibbs Sampling from Human
Feedback (M-GSHF) because the optimal policy is a layer-wise Gibbs distribution that generalizes
the result in Xiong et al. (2024). We now discuss some features of the framework as follows.

Reference model choice for controlling regular-
ization level. We unify the two different learn-
ing targets in equation 9 and equation 10 by
taking the reference model choice as a hyper-
parameter. First, if we fix the reference model
as the initial policy, i.e., πt,ref = π0, ∀t ∈ [T ],
we always search the optimal policy within
the KL ball centered at π0, and thus optimize
the KL-regularized target. In contrast, inspired
by the mirror descent (Nemirovskij & Yudin,
1983), if we update the reference policy every
iteration to be the policy learned in the last iter-
ation, i.e., πt,ref = π1
t−1, ∀t ∈ [T ], the cumula-
tive update can make the model to move away from the original π0 (while a constraint is made on
the per-iteration update magnitude) and we thus optimize the non-regularized target in equation 10.
See Figure 2 for an illustration.

Figure 2: Illustration of the difference between the
two learning objectives. Left: the KL-regularized
target as we do not update the reference model.
Right: the non-regularized target.

Non-symmetric policy choice for exploration-exploitation trade-off. We update our behavior policies
in a non-symmetric way. The first agent aims to extract the historical information we have gathered

5

Under review as a conference paper at ICLR 2025

so far and runs the M-DPO or M-DKO presented in Section 2.2. However, it is widely recognized
in RL studies (Sutton & 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. Therefore, the second agent will
strategically incorporate the uncertainty of the future relative to π1
t , which is referred
to as the exploration policy.

t to choose π2

A comprehensive theoretical analysis is derived for Algorithm 1, deferred to Appendix D due to
space constraint, with a focus on the KL-regularized target. Here we highlight the following in-
formal result (see Theorem 2 for the complete version), emphasizing the efficiency of Algorithm 1
guaranteed by a sublinear regret. 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 have
been developed in Cai et al. (2020).
Theorem 1 (Informal). Under the realizability assumption, with the KL-regularized target, the the-
oretical version of Algorithm 1 leads to a regret (defined in equation 13) that is sublinear in horizon
T for a broad class of reward and transition models.

The main take-away message from the theorem is that if we choose suitable exploration policy, the
online iterative learning is provably efficient. We also remark that without explicit mechanism to
encourage exploration, the randomness of the LLM itself is not sufficient to learn the optimal policy
(Zhang, 2022) if we do not make additional assumption.

Moving toward practical algorithm designs, the exploration is generally interpreted as increasing
the diversity of the collected data by adopting inference-time methods with the base DPO policy π1
t .
For instance, one may tune the sampling temperature as in Llama project (Touvron et al., 2023) or
use best-of-n sampling (Xu et al., 2023; Hoang Tran, 2024; Dong et al., 2024), where these methods
outperform the vanilla on-policy sampling with considerable margin. In this work, we mainly enrich
the generated data by various intermediate checkpoints, as done in the Claude project (Bai et al.,
2022). We refer this approach as mixture sampling. It is also natural to adopt reward-guided Monte
Carlo tree search (MCTS) (Xie et al., 2024b), which we leave for future work.

Algorithm 1 Online Iterative M-GSHF
1: Input: KL coefficient η > 0, horizon T > 0, initial policy π0, batch size m > 0.
2: Initialize D ← ∅ and π1
1 = π1,ref ← π0.
3: for t = 1, 2, · · · , T do
4:

1 = π2

t , τ 2 ∼ π2

t (e.g., using the M-DPO loss in equation 8

Sample m pairs (x, τ 1, τ 2, z) as Dt by x ∼ d0, τ 1 ∼ π1
t , receive the m preference signals z
following the Bradley-Terry model from Definition 1 and update the preference dataset D ← D ∪ Dt.
▷ Extract the empirically optimal policy from historical data
Practical: Perform the planning algorithms on D to get π1
or the M-KTO loss in equation 11)
Theoretical: Perform MLE on D to obtain model estimation ˆMt = (ˆut, ˆPt) as in equation 14 and
equation 15; call Oracle 2 with ˆMt, η, πt,ref to get π1
t
▷ Select the exploration policy to facilitate learning
Practical: Given π1
pling, inference parameters tuning and west-of-n sampling.
Theoretical: Given π1
t , choose π2
▷ Choose the reference model to control regularization level
Update πt+1,ref ← π1
ering the KL-regularized target

t when considering the non-regularized target; keep πt+1,ref ← π0 when consid-

t as an exploration policy using heuristic methods (such as mixture sam-

t as an exploration policy following equation 16

t , select π2

5:
6:

7:

8:
9:

10:
11:
12:

13: end for
14: Output: the best model in π1

1:T by a validation set.

3 EXPERIMENTS

3.1 EXPERIMENT SETUP

Task, datasets, and models. 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. To con-
struct the training prompt set, we use the prompts from MetaMathQA (Yu et al., 2023) and MMIQC

6

270
271

272
273
274
275
276
277

278
279
280
281
282
283

284
285
286
287
288

289
290
291
292
293
294

295
296
297
298
299
300

301
302
303
304
305
306

307
308
309
310
311

312
313
314
315
316
317

318
319
320
321
322
323

Under review as a conference paper at ICLR 2025

324
325

326
327
328
329
330
331

332
333
334
335
336
337

338
339
340
341
342

343
344
345
346
347
348

349
350
351
352
353
354

355
356
357
358
359
360

361
362
363
364
365

366
367
368
369
370
371

372
373
374
375
376
377

(Liu & Yao, 2024), which is an augmented prompt set from the 7.5K training problems of MATH
and 7.47K training problems of GSM8K. We provide an example of the data sample in Figure 1. 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.32 (Jiang et al., 2023), and Gemma2-it-9B. We first fine-tune
the model using a subset of the Open-MathInstruct dataset. The details of the SFT process are
provided in Appendix B.

Implementation of Iterative M-DPO and M-KTO. We run the iterative training for 3 epochs
in total. For each iteration, we have a prompt set of 20K questions and generate 20 responses per
prompt with current DPO model and 10 responses per prompt with the model from last iteration.
We check the final answer of these responses to determine their correctness. Then, for each prompt,
we randomly sample two responses with correct and incorrect final answers and add them into the
training samples. Then, we train the model on the collected samples using the M-DPO/M-KTO
loss. We also include an ablation of reference model choice. 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. The hyper-parameters are of M-KTO are mostly the same as the M-DPO. We also set
the λ+ = λ− = 1 following the original KTO paper (Ethayarajh et al., 2024).

Baselines. The existing literature mainly focuses on the synthetic data generation and SFT to teach
the models to use the external tool. 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
(Touvron et al., 2023). Another baseline is the single-turn online iterative DPO and KTO (Rafailov
et al., 2023; Ethayarajh et al., 2024), which ignore the problem structure (i.e., the external messages)
and treat the trajectory as a whole. In implementation, it means that we do not mask the tokens of
external messages.

3.2 MAIN RESULTS

We evaluate the models in the zero-shot setting and report the main results in Table 1. 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. Across all four base
models, iterative training with M-DPO or M-KTO consistently leads to notable improvements over
the initial SFT checkpoint 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, respec-
tively, 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 after SFT.

Iterative M-DPO and M-KTO surpass existing RLHF baselines. We also observe that the it-
erative M-DPO and M-KTO surpass other existing RLHF baselines. First, they consistently and
significantly outperform RAFT across all four base models. This is because RAFT only imitates
the correct trajectories, while the DPO-based and KTO-based algorithms further use the negative
signal from incorrect trajectories. We note that the SFT stage in our pipeline can also be viewed
as an application of RAFT. Consequently, our results should be interpreted to be that after the first

2We use the pre-trained version because the chat template of its instruct model from huggingface is not

consistent with their own codebase.

7

Under review as a conference paper at ICLR 2025

378
379

380
381
382
383
384
385

386
387
388
389
390
391

392
393
394
395
396

397
398
399
400
401
402

403
404
405
406
407
408

409
410
411
412
413
414

415
416
417
418
419

420
421
422
423
424
425

426
427
428
429
430
431

Table 1: Main results of different methods on the test sets of GSM8K and MATH. †: the model
serves as the starting checkpoint of other methods. The results of the CoT methods are borrowed
from the technical reports (Toshniwal et al., 2024; Gou et al., 2023b). For iterative M-DPO/M-
KTO, we update the reference model by default if not specified. The gains relative to the SFT
starting checkpoint are marked by ↑.

Base Model

WizardMath-7B
WizardMath-13B
WizardMath-70B

CodeLLaMA-2-7B
CodeLLaMA-2-13B
CodeLLaMA-2-34B
CodeLLaMA-2-70B

Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B

CodeGemma-1.1-it-7B
CodeGemma-1.1-it-7B
CodeGemma-1.1-it-7B
CodeGemma-1.1-it-7B
CodeGemma-1.1-it-7B
CodeGemma-1.1-it-7B

Mistral-7B-v0.3
Mistral-7B-v0.3
Mistral-7B-v0.3
Mistral-7B-v0.3
Mistral-7B-v0.3
Mistral-7B-v0.3

Gemma-2-it-9B
Gemma-2-it-9B
Gemma-2-it-9B
Gemma-2-it-9B
Gemma-2-it-9B
Gemma-2-it-9B

Method

SFT for CoT
SFT for CoT
SFT for CoT

SFT
SFT
SFT
SFT

SFT†
RAFT
Iterative Single-turn DPO
Iterative Single-turn KTO
Iterative M-DPO + fixed reference
M-DPO Iteration 1
M-DPO Iteration 2
M-DPO Iteration 3
Iterative M-KTO

SFT†
RAFT
Iterative Single-turn DPO
Iterative Single-turn KTO
Iterative M-DPO
Iterative M-KTO

SFT†
RAFT
Iterative Single-turn DPO
Iterative Single-turn KTO
Iterative M-DPO
Iterative M-KTO

SFT†
RAFT
Iterative Single-turn DPO
Iterative Single-turn KTO
Iterative M-DPO
Iterative M-KTO

with Tool

GSM8K

MATH

AVG

✗
✗
✗

✓
✓
✓
✓

✓
✓
✓
✓
✓
✓
✓
✓
✓

✓
✓
✓
✓
✓
✓

✓
✓
✓
✓
✓
✓

✓
✓
✓
✓
✓
✓

54.9
63.9
81.6

75.9
78.8
80.7
84.6

77.5
79.2
81.7
80.6
79.9
81.5
82.5
83.9 ↑6.4
82.1 ↑4.6

77.3
78.8
79.1
80.2
81.5 ↑4.2
81.6 ↑4.3

77.8
79.8
79.8
81.3
82.3 ↑4.5
81.7 ↑3.9

84.1
84.2
85.2
85.4
86.3 ↑2.2
86.1 ↑2.0

10.7
14.0
22.7

43.6
45.5
48.3
50.7

46.1
47.3
48.9
49.0
48.0
49.1
49.7
51.2 ↑5.1
49.5 ↑3.4

46.4
48.4
48.9
48.6
50.1 ↑3.7
49.6 ↑3.2

42.7
43.7
45.1
46.3
47.5 ↑4.8
46.7 ↑4.0

51.0
52.6
53.1
52.9
54.5 ↑3.5
54.5 ↑3.5

32.8
39.0
52.2

59.8
62.2
64.5
67.7

61.8
63.3
65.3
64.8
64.0
65.3
66.1
67.6 ↑5.8
65.8 ↑4.0

61.9
63.6
64.0
64.4
65.8 ↑4.0
65.6 ↑3.8

60.3
61.8
62.5
63.8
64.9 ↑4.7
64.2 ↑4.0

67.6
68.4
69.2
69.2
70.4 ↑2.9
70.3 ↑2.8

stage of SFT, algorithms with negative signal are more sample efficient. Moreover, while the on-
line iterative single-turn DPO (KTO) also gives a better 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. We also
present a representative example we encounter in Figure 5, where LLMs generate poorly constructed
code resulting in anomalous and lengthy external messages. Forcing LLMs to learn to predict these
messages can significantly hurt the model’s reasoning abilities.

Iterative training and reference update lead to better performance. Using Gemma-1.1-it-7B
with M-DPO as an example, we observe that online iterative training leads to better results. The
GSM8K test accuracy increases 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 aligns with our theoretical insight that iterative training helps models
progressively explore and learn the optimal policy. Additionally, if the reference model remains
fixed at the SFT policy, the final performance is notably worse compared to updating the refer-
ence model at each iteration. This likely occurs because the algorithm, in this case, optimizes the
non-regularized reward, and the rewards in mathematical reasoning tasks are more accurate than in
general chat tasks, leading to better in-domain performance. A detailed ablation study on the impact
of KL regularization is deferred to the next section.

8

Under review as a conference paper at ICLR 2025

Figure 3: 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 3. A question is solved if at
least one of the n sampled trajectories is correct. We find that preference learning improves pass@n
accuracy only when n is small. For n > 16, all models perform similarly on GSM8K and MATH,
indicating that iterative M-DPO does not introduce new knowledge but instead enhances the quality
of top-n responses. This observation also aligns with the result of CoT reasoning (Shao et al., 2024).

3.3 ABLATION STUDY AND DISCUSSION

GSM8K MATH

Method

SFT
update reference + η = 0.01
update reference + η = 0.1
update reference + η = 0.5

Moderate KL regularization balances per-
iteration improvement and exploration.
The effectiveness of iterative DPO is highly
dependent on the reference model and KL
coefficient.
In our ablation study, we first
consider two different choices of the reference
model: (1) using the fixed reference model π0;
(2) updating the reference model to the last
iteration’s model at each round, which can be
viewed as a trade-off between the generation
diversity and reward optimization. As shown
in Table 3.3, models with an updated reference
model outperform those with a fixed reference model. We hypothesize that in reasoning tasks, the
correct reasoning paths are highly concentrated, making diversity less crucial so optimizing the
non-regularized reward gives superior model performance.

Table 2: Ablation study of the impact of KL regu-
larization on iterative M-DPO.

fixed reference + η = 0.1

46.1
50.1
51.2
49.7

77.5
81.7
83.9
82.8

48.0

79.9

Previous work (Tunstall et al., 2023) on offline DPO suggests that a lower KL coefficient (0.01)
improves performance by allowing the model to deviate more from the SFT model π0.
In our
ablation study, we search the KL coefficient η ∈ {0.01, 0.1, 0.5}. According to Table 3.3, we find
that the strongest model is obtained by a moderate KL coefficient of 0.1, outperforming both 0.01
and 0.5. To explain this, we plot the GSM8K test accuracy (Figure 4) during iterative training. In the
first iteration, lower KL values show larger improvements, consistent with Tunstall et al. (2023)’s
results. However, models trained with very low KL coefficients lose diversity quickly, reducing their
ability to generate diverse trajectories for later training, leading to diminishing returns in subsequent
iterations. Conversely, a higher KL coefficient of 0.5 imposes too much regularization, limiting
improvement per iteration. In summary, for online iterative training, a balance between per-iteration
improvement and exploration efficiency is key to optimizing overall performance, an intuition that
also applies to sampling strategies and other experimental techniques.

The impact of sampling strategy: data diversity and coverage are crucial. During iterative
training of Gemma-1.1-it-7B, we see an increase in correct trajectories from 47% in the first iter-
ation to 76% in last iteration. Moreover, as the reference model updates at each step, trajectory
diversity declines, which is critical for DPO/KTO training due to its contrastive nature. We follow
Bai et al. (2022); Dong et al. (2024) to explore two data collection strategies: (1) on-policy sam-

9

432
433

434
435
436
437
438
439

440
441
442
443
444
445

446
447
448
449
450

451
452
453
454
455
456

457
458
459
460
461
462

463
464
465
466
467
468

469
470
471
472
473

474
475
476
477
478
479

480
481
482
483
484
485

1248163264n: the number of candidates80859095Pass@n (%)GSM8K Test AccuracyGemma-7B SFTGemma-7B M-DPO + Fixed ReferenceGemma-7B M-DPO + Update Reference1248163264n: the number of candidates405060708090Pass@n (%)MATH Test AccuracyGemma-7B SFTGemma-7B M-DPO + Fixed ReferenceGemma-7B M-DPO + Update ReferenceUnder review as a conference paper at ICLR 2025

Figure 4: Left:
Right: the test accuracy on MATH dataset with different sampling strategies.

the test accuracy on GSM8K dataset with different levels of KL regularization.

pling (trajectories sampled from the current model) and (2) mixture sampling (20 trajectories from
the current model and 10 from model of previous iteration). As shown in Table 6, mixture sampling
significantly outperforms on-policy sampling, particularly in the third iteration where on-policy sam-
pling fails to improve MATH test accuracy. This highlights the importance of diversity in iterative
training and aligns with previous findings that advanced exploration strategies help prevent diver-
sity collapse and improve preference learning (Bai et al., 2022; Touvron et al., 2023; Xiong et al.,
2024; Pace et al., 2024; Dong et al., 2024). It would also be interested to explore more advanced
exploration strategy like MCTS in the future study.

Method

GSM8K MATH

To ensure both correct and incorrect reason-
ing paths exist, we collected N trajectories per
prompt. A larger N generally improves prompt
coverage, as more samples are needed for diffi-
cult problems. For example, in iteration 1, with
N=30, 92.5% of the prompts are covered, com-
pared to 83.0% for N=12 and 60% for N=6. See
Figure 3 for an illustration of the relationship
between pass@1 and N. However, increasing N
also increases computational costs. In our abla-
tion study (Table 3.3), we find that increasing N
from 6 to 12 leads to a significant performance boost, reflecting better coverage for complex prob-
lems. However, increasing N from 12 to 30 yields only minor improvements, suggesting that the
benefits of larger N diminish quickly in vanilla rejection sampling. We expect that difficulty-aware
sampling can lead to a better performance, while maintaining a moderate inference cost.

Table 3: Ablation study of the sampling strategy
with iterative M-DPO and Gemma-1.1-it-7B.

SFT
N=30 + Mixture
N=12 + Mixture
N=6 + Mixture
N=30 + On-policy

77.5
83.9
83.5
82.0
83.1

46.1
51.2
51.2
49.2
49.5

4 CONCLUSION, LIMITATION, AND FUTURE RESEARCH DIRECTION

In this paper, we demonstrate that preference learning, as an alternative to supervised fine-tuning,
further enhances the performance of tool-integrated reasoning LLMs after SFT. We introduce an
online iterative multi-turn direct preference optimization algorithm, validated through extensive ex-
periments across multiple base models. Results show significant improvements in pass@1 over the
SFT policy, particularly on benchmarks like GSM8K and MATH. Ablation studies highlight the im-
portance of balancing per-iteration improvement with exploration, which is achieved by moderate
levels of KL regularization and strategic exploration choices.

Several avenues for improvement remain unexplored. Our current approach only uses final result
checks as preference signals, limiting the comparison between trajectories with correct or incorrect
answers. One may use step-wise reward signal (Lightman et al., 2023) in the data ranking stage.
Meanwhile, the fine-grained reward signals could enable the use of advanced exploration strategies
like west-of-n sampling (Pace et al., 2024), or MCTS (Xie et al., 2024b) in our heuristic exploration
implementation. 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 oppo-
nents. 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.

10

486
487

488
489
490
491
492
493

494
495
496
497
498
499

500
501
502
503
504

505
506
507
508
509
510

511
512
513
514
515
516

517
518
519
520
521
522

523
524
525
526
527

528
529
530
531
532
533

534
535
536
537
538
539

0123Iteration767778798081828384AccuracyGSM8K Test AccuracySFTeta = 0.01eta = 0.1eta = 0.5eta = 0.1 + fixed reference0123Iteration45464748495051AccuracyMATH Test AccuracySFT 3 epochsMixture SamplingOn-policy SamplingUnder review as a conference paper at ICLR 2025

540
541

542
543
544
545
546
547

548
549
550
551
552
553

554
555
556
557
558

559
560
561
562
563
564

565
566
567
568
569
570

571
572
573
574
575
576

577
578
579
580
581

582
583
584
585
586
587

588
589
590
591
592
593

REPRODUCIBILITY STATEMENT

We believe that making the result reproducible is important. Following the author guidance of
ICLR, we present a reproducibility statement here to help the interested readers to reproduce our
result. Most implementation details, including hyperparameters, are provided in Section 3.1 and
Appendix B. Additionally, we have open-sourced our training code along with a step-by-step guide,
using Gemma-1.1-it-7B as an example. We have also made the processed SFT dataset, prompt
set, and the training data for the first iteration of M-DPO/M-KTO available for easy download (see
supplemental materials for details). The RLHF experiments of this paper are run with 8xA100
80G GPUs, where an additional machine with 8xA100 40G GPUs is also used to accelerate data
collection and model evaluation. The main experiment of this paper can be reproduced within 24
- 48 hours with this setup. To improve the readability of this paper, we provide a notation table in
Appendix A. The informal version of our main theoretical result is summarized in Theorem 1 is
re-stated in Theorem 2 and its proof is provided in Appendix D.

REFERENCES

Qwen2 technical report. 2024.

Alekh Agarwal, Yujia Jin, and Tong Zhang. VOQL: Towards optimal regret in model-free rl with
nonlinear function approximation. In The Thirty Sixth Annual Conference on Learning Theory,
pp. 987–1063. PMLR, 2023.

Anthropic.

Introducing claude.

2023. URL https://www.anthropic.com/index/

introducing-claude.

Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit

problem. Machine learning, 47:235–256, 2002.

Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal
Valko, and R´emi Munos. A general theoretical paradigm to understand learning from human
preferences. arXiv preprint arXiv:2310.12036, 2023.

Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn
Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training a helpful and harmless
assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862,
2022.

Qi Cai, Zhuoran Yang, Chi Jin, and Zhaoran Wang. Provably efficient exploration in policy opti-

mization. In International Conference on Machine Learning, pp. 1283–1294. PMLR, 2020.

Shicong Cen, Jincheng Mei, Katayoon Goshvadi, Hanjun Dai, Tong Yang, Sherry Yang, Dale Schu-
urmans, Yuejie Chi, and Bo Dai. Value-incentivized preference optimization: A unified approach
to online and offline rlhf. arXiv preprint arXiv:2405.19320, 2024.

Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Step-level value preference optimization

for mathematical reasoning. arXiv preprint arXiv:2406.10858, 2024a.

Wenhu Chen, Xueguang Ma, Xinyi Wang, and William W Cohen. Program of thoughts prompt-
ing: Disentangling computation from reasoning for numerical reasoning tasks. arXiv preprint
arXiv:2211.12588, 2022.

Zixiang Chen, Yihe Deng, Huizhuo Yuan, Kaixuan Ji, and Quanquan Gu. Self-play fine-tuning
converts weak language models to strong language models. arXiv preprint arXiv:2401.01335,
2024b.

Leshem Choshen, Lior Fox, Zohar Aizenbud, and Omri Abend. On the weaknesses of reinforcement

learning for neural machine translation. arXiv preprint arXiv:1907.01752, 2019.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser,
Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John
Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168,
2021a.

11

Under review as a conference paper at ICLR 2025

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser,
Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to
solve math word problems. arXiv preprint arXiv:2110.14168, 2021b.

Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao,
Jipeng Zhang, KaShun SHUM, and Tong Zhang. RAFT: Reward ranked finetuning for generative
foundation model alignment. Transactions on Machine Learning Research, 2023. ISSN 2835-
8856. URL https://openreview.net/forum?id=m7p5O7zblY.

Hanze Dong, Wei Xiong, Bo Pang, Haoxiang Wang, Han Zhao, Yingbo Zhou, Nan Jiang, Doyen
Sahoo, Caiming Xiong, and Tong Zhang. Rlhf workflow: From reward modeling to online rlhf.
arXiv preprint arXiv:2405.07863, 2024.

Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry
Rudolph, and Aleksander Madry. Implementation matters in deep policy gradients: A case study
on ppo and trpo. arXiv preprint arXiv:2005.12729, 2020.

Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model

alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01306, 2024.

Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and
Graham Neubig. Pal: Program-aided language models. In International Conference on Machine
Learning, pp. 10764–10799. PMLR, 2023.

Team Gemini, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu,
Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly
capable multimodal models. arXiv preprint arXiv:2312.11805, 2023.

Zhibin Gou, Zhihong Shao, Yeyun Gong, Yelong Shen, Yujiu Yang, Nan Duan, and Weizhu Chen.
Critic: Large language models can self-correct with tool-interactive critiquing. arXiv preprint
arXiv:2305.11738, 2023a.

Zhibin Gou, Zhihong Shao, Yeyun Gong, Yujiu Yang, Minlie Huang, Nan Duan, Weizhu Chen,
et al. Tora: A tool-integrated reasoning agent for mathematical problem solving. arXiv preprint
arXiv:2309.17452, 2023b.

Lin Gui, Cristina Gˆarbacea, and Victor Veitch. Bonbon alignment for large language models and

the sweetness of best-of-n sampling. arXiv preprint arXiv:2406.00832, 2024.

Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek
Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, et al. Reinforced self-training
(rest) for language modeling. arXiv preprint arXiv:2308.08998, 2023.

Shangmin Guo, Wei Xiong, and Chaoqi Wang. ”alignment guidebook. Notion Blog, 2024a.

Shangmin Guo, Biao Zhang, Tianlin Liu, Tianqi Liu, Misha Khalman, Felipe Llinares, Alexandre
Rame, Thomas Mesnard, Yao Zhao, Bilal Piot, et al. Direct language model alignment from
online ai feedback. arXiv preprint arXiv:2402.04792, 2024b.

Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song,
and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv
preprint arXiv:2103.03874, 2021.

Braden Hancock Hoang Tran, Chris Glaze.

Snorkel-mistral-pairrm-dpo.

//huggingface.co/snorkelai/Snorkel-Mistral-PairRM-DPO, 2024.
https://huggingface.co/snorkelai/Snorkel-Mistral-PairRM-DPO.

https:
URL

Jiwoo Hong, Noah Lee, and James Thorne. Orpo: Monolithic preference optimization without

reference model. arXiv preprint arXiv:2403.07691, 2(4):5, 2024.

Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot,
Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al.
Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.

12

594
595

596
597
598
599
600
601

602
603
604
605
606
607

608
609
610
611
612

613
614
615
616
617
618

619
620
621
622
623
624

625
626
627
628
629
630

631
632
633
634
635

636
637
638
639
640
641

642
643
644
645
646
647

Under review as a conference paper at ICLR 2025

648
649

650
651
652
653
654
655

656
657
658
659
660
661

662
663
664
665
666

667
668
669
670
671
672

673
674
675
676
677
678

679
680
681
682
683
684

685
686
687
688
689

690
691
692
693
694
695

696
697
698
699
700
701

Fangkai Jiao, Chengwei Qin, Zhengyuan Liu, Nancy F Chen, and Shafiq Joty. Learning planning-
arXiv preprint

based reasoning by trajectories collection and process reward synthesizing.
arXiv:2402.00658, 2024.

Xin Lai, Zhuotao Tian, Yukang Chen, Senqiao Yang, Xiangru Peng, and Jiaya Jia. Step-dpo: Step-
wise preference optimization for long-chain reasoning of llms. arXiv preprint arXiv:2406.18629,
2024.

Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan
Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let’s verify step by step. arXiv preprint
arXiv:2305.20050, 2023.

Yong Lin, Lu Tan, Hangyu Lin, Zeming Zheng, Renjie Pi, Jipeng Zhang, Shizhe Diao, Haoxiang
Wang, Han Zhao, Yuan Yao, et al. Speciality vs generality: An empirical study on catastrophic
forgetting in fine-tuning foundation models. arXiv preprint arXiv:2309.06256, 2023.

Haoxiong Liu and Andrew Chi-Chih Yao. Augmenting math word problems via iterative question

composing. arXiv preprint arXiv:2401.09003, 2024.

Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle: A generic model-
based algorithm for partially observable sequential decision making. In Proceedings of the 55th
Annual ACM Symposium on Theory of Computing, pp. 363–376, 2023a.

Tianqi Liu, Yao Zhao, Rishabh Joshi, Misha Khalman, Mohammad Saleh, Peter J Liu, and
arXiv preprint

Jialu Liu. Statistical rejection sampling improves preference optimization.
arXiv:2309.06657, 2023b.

Tianqi Liu, Zhen Qin, Junru Wu, Jiaming Shen, Misha Khalman, Rishabh Joshi, Yao Zhao, Moham-
mad Saleh, Simon Baumgartner, Jialu Liu, et al. Lipo: Listwise preference optimization through
learning-to-rank. arXiv preprint arXiv:2402.01878, 2024a.

Zhihan Liu, Miao Lu, Shenao Zhang, Boyi Liu, Hongyi Guo, Yingxiang Yang, Jose Blanchet, and
Zhaoran Wang. Provably mitigating overoptimization in rlhf: Your sft loss is implicitly an adver-
sarial regularizer. arXiv preprint arXiv:2405.16436, 2024b.

Zimu Lu, Aojun Zhou, Ke Wang, Houxing Ren, Weikang Shi, Junting Pan, and Mingjie Zhan. Step-
controlled dpo: Leveraging stepwise error for enhanced mathematical reasoning. arXiv preprint
arXiv:2407.00782, 2024.

Haipeng Luo, Qingfeng Sun, Can Xu, Pu Zhao, Jianguang Lou, Chongyang Tao, Xiubo Geng, Qing-
wei Lin, Shifeng Chen, and Dongmei Zhang. Wizardmath: Empowering mathematical reasoning
for large language models via reinforced evol-instruct. arXiv preprint arXiv:2308.09583, 2023.

Yu Meng, Mengzhou Xia, and Danqi Chen. Simpo: Simple preference optimization with a

reference-free reward. arXiv preprint arXiv:2405.14734, 2024.

Meta. Introducing meta llama 3: The most capable openly available llm to date. Meta AI Blog,

2024. https://ai.meta.com/blog/meta-llama-3/.

Swaroop Mishra, Matthew Finlayson, Pan Lu, Leonard Tang, Sean Welleck, Chitta Baral, Tanmay
Rajpurohit, Oyvind Tafjord, Ashish Sabharwal, Peter Clark, et al. Lila: A unified benchmark for
mathematical reasoning. arXiv preprint arXiv:2210.17517, 2022.

Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orca-math: Unlocking

the potential of slms in grade school math. arXiv preprint arXiv:2402.14830, 2024.

R´emi Munos, Michal Valko, Daniele Calandriello, Mohammad Gheshlaghi Azar, Mark Rowland,
Zhaohan Daniel Guo, Yunhao Tang, Matthieu Geist, Thomas Mesnard, Andrea Michi, et al. Nash
learning from human feedback. arXiv preprint arXiv:2312.00886, 2023.

Arkadij Semenoviˇc Nemirovskij and David Borisovich Yudin. Problem complexity and method

efficiency in optimization. 1983.

OpenAI. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.

13

Under review as a conference paper at ICLR 2025

702
703

704
705
706
707
708
709

710
711
712
713
714
715

716
717
718
719
720

721
722
723
724
725
726

727
728
729
730
731
732

733
734
735
736
737
738

739
740
741
742
743

744
745
746
747
748
749

750
751
752
753
754
755

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong
Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow
instructions with human feedback. Advances in Neural Information Processing Systems, 35:
27730–27744, 2022.

Aliz´ee Pace, Jonathan Mallinson, Eric Malmi, Sebastian Krause, and Aliaksei Severyn. West-of-n:
Synthetic preference generation for improved reward modeling. arXiv preprint arXiv:2401.12086,
2024.

Richard Yuanzhe Pang, Weizhe Yuan, Kyunghyun Cho, He He, Sainbayar Sukhbaatar, and Jason
Weston. Iterative reasoning preference optimization. arXiv preprint arXiv:2404.19733, 2024.

Renjie Pi, Tianyang Han, Wei Xiong, Jipeng Zhang, Runtao Liu, Rui Pan, and Tong Zhang.
Strengthening multimodal large language model with bootstrapped preference optimization.
arXiv preprint arXiv:2403.08730, 2024.

Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea
Finn. Direct preference optimization: Your language model is secretly a reward model. arXiv
preprint arXiv:2305.18290, 2023.

Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From r to q*: Your language model is

secretly a q-function. arXiv preprint arXiv:2404.12358, 2024.

Corby Rosset, Ching-An Cheng, Arindam Mitra, Michael Santacroce, Ahmed Awadallah, and
Tengyang Xie. Direct nash optimization: Teaching language models to self-improve with general
preferences. arXiv preprint arXiv:2404.03715, 2024.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy

optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.

Lior Shani, Aviv Rosenberg, Asaf Cassel, Oran Lang, Daniele Calandriello, Avital Zipori, Hila
Noga, Orgad Keller, Bilal Piot, Idan Szpektor, et al. Multi-turn reinforcement learning from
preference human feedback. arXiv preprint arXiv:2405.14655, 2024.

Zhihong Shao, Fei Huang, and Minlie Huang. Chaining simultaneous thoughts for numerical rea-

soning. arXiv preprint arXiv:2211.16482, 2022.

Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu,
and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language
models. arXiv preprint arXiv:2402.03300, 2024.

Avi Singh, John D Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Peter J Liu, James
Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, et al. Beyond human data: Scaling self-training
for problem-solving with language models. arXiv preprint arXiv:2312.06585, 2023.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.

Gokul Swamy, Christoph Dann, Rahul Kidambi, Zhiwei Steven Wu, and Alekh Agarwal. A
arXiv preprint

minimaximalist approach to reinforcement learning from human feedback.
arXiv:2401.04056, 2024.

Fahim Tajwar, Anikait Singh, Archit Sharma, Rafael Rafailov, Jeff Schneider, Tengyang Xie, Ste-
fano Ermon, Chelsea Finn, and Aviral Kumar. Preference fine-tuning of llms should leverage
suboptimal, on-policy data. arXiv preprint arXiv:2404.14367, 2024.

Yunhao Tang, Zhaohan Daniel Guo, Zeyu Zheng, Daniele Calandriello, R´emi Munos, Mark Row-
land, Pierre Harvey Richemond, Michal Valko, Bernardo ´Avila Pires, and Bilal Piot. Gen-
arXiv preprint
eralized preference optimization: A unified approach to offline alignment.
arXiv:2402.05749, 2024.

CodeGemma Team.

Codegemma: Open code models based on gemma.

arXiv preprint

arXiv:2406.11409, 2024.

14

Under review as a conference paper at ICLR 2025

756
757

758
759
760
761
762
763

764
765
766
767
768
769

770
771
772
773
774

775
776
777
778
779
780

781
782
783
784
785
786

787
788
789
790
791
792

793
794
795
796
797

798
799
800
801
802
803

804
805
806
807
808
809

Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya
Pathak, Laurent Sifre, Morgane Rivi`ere, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open
models based on gemini research and technology. arXiv preprint arXiv:2403.08295, 2024.

Yuxuan Tong, Xiwen Zhang, Rui Wang, Ruidong Wu, and Junxian He. Dart-math: Difficulty-aware

rejection tuning for mathematical problem-solving. 2024.

Shubham Toshniwal, Ivan Moshkov, Sean Narenthiran, Daria Gitman, Fei Jia, and Igor Git-
arXiv preprint

man. Openmathinstruct-1: A 1.8 million math instruction tuning dataset.
arXiv:2402.10176, 2024.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko-
lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda-
tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.

Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada,
Shengyi Huang, Leandro von Werra, Cl´ementine Fourrier, Nathan Habib, et al. Zephyr: Direct
distillation of lm alignment. arXiv preprint arXiv:2310.16944, 2023.

Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia
Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process-and
outcome-based feedback. arXiv preprint arXiv:2211.14275, 2022.

Peiyi Wang, Lei Li, Zhihong Shao, RX Xu, Damai Dai, Yifei Li, Deli Chen, Y Wu, and Zhifang
Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations. CoRR,
abs/2312.08935, 2023a.

Xingyao Wang, Zihan Wang, Jiateng Liu, Yangyi Chen, Lifan Yuan, Hao Peng, and Heng Ji. Mint:
Multi-turn interactive evaluation for tool-augmented llms with language feedback. In Proc. The
Twelfth International Conference on Learning Representations (ICLR2024), 2024.

Yuanhao Wang, Qinghua Liu, and Chi Jin. Is rlhf more difficult than standard rl? arXiv preprint

arXiv:2306.14111, 2023b.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny
Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in
neural information processing systems, 35:24824–24837, 2022.

Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement

learning. Machine learning, 8:229–256, 1992.

Ronald J Williams and Jing Peng. Function optimization using connectionist reinforcement learning

algorithms. Connection Science, 3(3):241–268, 1991.

Tengyang Xie, Dylan J Foster, Yu Bai, Nan Jiang, and Sham M Kakade. The role of coverage in

online reinforcement learning. arXiv preprint arXiv:2210.04157, 2022.

Tengyang Xie, Dylan J Foster, Akshay Krishnamurthy, Corby Rosset, Ahmed Awadallah, and
Alexander Rakhlin. Exploratory preference optimization: Harnessing implicit q*-approximation
for sample-efficient rlhf. arXiv preprint arXiv:2405.21046, 2024a.

Yuxi Xie, Anirudh Goyal, Wenyue Zheng, Min-Yen Kan, Timothy P Lillicrap, Kenji Kawaguchi,
and Michael Shieh. Monte carlo tree search boosts reasoning via iterative preference learning.
arXiv preprint arXiv:2405.00451, 2024b.

Wei Xiong, Hanze Dong, Chenlu Ye, Ziqi Wang, Han Zhong, Heng Ji, Nan Jiang, and Tong Zhang.
Iterative preference learning from human feedback: Bridging theory and practice for rlhf under
kl-constraint. In Forty-first International Conference on Machine Learning, 2024.

Jing Xu, Andrew Lee, Sainbayar Sukhbaatar, and Jason Weston. Some things are more cringe than
others: Preference optimization with the pairwise cringe loss. arXiv preprint arXiv:2312.16682,
2023.

15

Under review as a conference paper at ICLR 2025

810
811

812
813
814
815
816
817

818
819
820
821
822
823

824
825
826
827
828

829
830
831
832
833
834

835
836
837
838
839
840

841
842
843
844
845
846

847
848
849
850
851

852
853
854
855
856
857

858
859
860
861
862
863

Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao.
React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03629,
2022.

Chenlu Ye, Wei Xiong, Yuheng Zhang, Nan Jiang, and Tong Zhang. A theoretical analysis of
nash learning from human feedback under general kl-regularized preference. arXiv preprint
arXiv:2402.07314, 2024.

Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhen-
guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions
for large language models. arXiv preprint arXiv:2309.12284, 2023.

Lifan Yuan, Ganqu Cui, Hanbin Wang, Ning Ding, Xingyao Wang, Jia Deng, Boji Shan, Huimin
Chen, Ruobing Xie, Yankai Lin, et al. Advancing llm reasoning generalists with preference trees.
arXiv preprint arXiv:2404.02078, 2024.

Zheng Yuan, Hongyi Yuan, Chengpeng Li, Guanting Dong, Chuanqi Tan, and Chang Zhou. Scal-
ing relationship on learning mathematical reasoning with large language models. arXiv preprint
arXiv:2308.01825, 2023a.

Zheng Yuan, Hongyi Yuan, Chuanqi Tan, Wei Wang, Songfang Huang, and Fei Huang. Rrhf:
Rank responses to align language models with human feedback without tears. arXiv preprint
arXiv:2304.05302, 2023b.

Xiang Yue, Ge Zhang Xingwei Qu, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen.
Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint
arXiv:2309.05653, 2023.

Xiang Yue, Tuney Zheng, Ge Zhang, and Wenhu Chen. Mammoth2: Scaling instructions from the

web. arXiv preprint arXiv:2405.03548, 2024.

Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with

reasoning. Advances in Neural Information Processing Systems, 35:15476–15488, 2022.

Wenhao Zhan, Masatoshi Uehara, Nathan Kallus, Jason D Lee, and Wen Sun. Provable offline

reinforcement learning with human feedback. arXiv preprint arXiv:2305.14816, 2023.

Beichen Zhang, Kun Zhou, Xilin Wei, Xin Zhao, Jing Sha, Shijin Wang, and Ji-Rong Wen. Evalu-
ating and improving tool-augmented computation-intensive math reasoning. Advances in Neural
Information Processing Systems, 36, 2024a.

Shenao Zhang, Donghan Yu, Hiteshi Sharma, Ziyi Yang, Shuohang Wang, Hany Hassan, and Zhao-
ran Wang. Self-exploring language models: Active preference elicitation for online alignment.
arXiv preprint arXiv:2405.19332, 2024b.

Tong Zhang. Feel-good thompson sampling for contextual bandits and reinforcement learning. SIAM

Journal on Mathematics of Data Science, 4(2):834–857, 2022.

Tong Zhang. Mathematical analysis of machine learning algorithms. Cambridge University Press,

2023.

Yuheng Zhang, Dian Yu, Baolin Peng, Linfeng Song, Ye Tian, Mingyue Huo, Nan Jiang, Haitao
Mi, and Dong Yu. Iterative nash policy optimization: Aligning llms with general preferences via
no-regret learning. arXiv preprint arXiv:2407.00617, 2024c.

Yao Zhao, Rishabh Joshi, Tianqi Liu, Misha Khalman, Mohammad Saleh, and Peter J Liu. Slic-hf:
Sequence likelihood calibration with human feedback. arXiv preprint arXiv:2305.10425, 2023.

Chujie Zheng, Ziqi Wang, Heng Ji, Minlie Huang, and Nanyun Peng. Weak-to-strong extrapolation

expedites alignment. arXiv preprint arXiv:2404.16792, 2024.

Kunhao Zheng, Jesse Michael Han, and Stanislas Polu. Minif2f: a cross-system benchmark for

formal olympiad-level mathematics. arXiv preprint arXiv:2109.00110, 2021.

16

Under review as a conference paper at ICLR 2025

Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang, Zhuoran Yang, and Tong Zhang.
Gec: A unified framework for interactive decision making in mdp, pomdp, and beyond. arXiv
preprint arXiv:2211.01962, 2022.

Han Zhong, Guhao Feng, Wei Xiong, Li Zhao, Di He, Jiang Bian, and Liwei Wang. Dpo meets ppo:

Reinforced token optimization for rlhf. arXiv preprint arXiv:2404.18922, 2024.

Denny Zhou, Nathanael Sch¨arli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur-
mans, Claire Cui, Olivier Bousquet, Quoc Le, et al. Least-to-most prompting enables complex
reasoning in large language models. arXiv preprint arXiv:2205.10625, 2022.

Xinyu Zhu, Junjie Wang, Lin Zhang, Yuxiang Zhang, Yongfeng Huang, Ruyi Gan, Jiaxing Zhang,
and Yujiu Yang. Solving math word problems via cooperative reasoning induced language mod-
els. arXiv preprint arXiv:2210.16257, 2022.

Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal

entropy. Carnegie Mellon University, 2010.

17

864
865

866
867
868
869
870
871

872
873
874
875
876
877

878
879
880
881
882

883
884
885
886
887
888

889
890
891
892
893
894

895
896
897
898
899
900

901
902
903
904
905

906
907
908
909
910
911

912
913
914
915
916
917

Under review as a conference paper at ICLR 2025

A NOTATION TABLE, RELATED WORK, AND MISSING DETAILS

A.1 NOTATION TABLE

Notation
x, X
d0
sh ∈ S, ah ∈ A, oh
H
P∗ = [P∗

h=1

h=1

h=1

h=1

h=1

h]H
τ = (x, y)
u∗
M∗ = (S, A, H, P∗, d0, u∗)
σ(·)
z ∈ {0, 1}
π = [πh]H
M = (S, A, H, P, d0, u)
πref = [πref,h]H
J(π; M, πref )
η
QM = [QM,h]H
VM = [VM,h]H
πM = [πM,h]H
LM-DPO(·)
LM-KTO(·)
J(π)
π∗ = [π∗
h]H
t , π2
π1
t
Reg(T )
U, P
B
ˆut, ˆPt
(cid:101)Ut, (cid:101)Pt
c1, c2, c
κ
dU
dP , ξ(·)
TV(·, ·)

h=1

h=1

Description
The prompt and the prompt space.
The distribution of initial state (prompt).
The state, action, and observation.
Episode length, e.g., the maximal number of tool calls.
The true observation kernel.
τ is a trajectory and y is the completion part, i.e., we exclude x from τ .
The true utility function associated with the BT model defined in Definition 1.
The true model with observation kernel P∗ and utility function u∗
σ(z) = 1/(1 + exp(−z)) is the sigmoid function.
Preference signal.
The policy, which is parameterized by the LLM.
One arbitrary environment with observation kernel P and utility function u.
One arbitrary reference policy.
The KL-regularized target (equation 2) with environment M and reference πref .
The coefficient of KL penalty, defined in equation 2.
The optimal Q-values associated with J(π; M, πref ), defined in equation 3.
The optimal V -values associated with J(π; M, πref ), defined in equation 4.
The optimal policy associated with J(π; M, πref ), defined in equation 4.
M-DPO loss, defined in equation 8.
M-KTO loss, defined in equation 11.
The abbreviation of J(π; M∗, π0), defined in equation 12.
The optimal policy associated with J(π).
The main and exploration policy at round t
Regret over horizon T , defined in equation 13.
Known sets such that u∗ ∈ U and P∗ ∈ P
Assuming u∗(x, y) ∈ [0, B], ∀(x, y).
MLE of u∗ and P∗ at round t, defined in equation 14 and equation 15.
Confidences sets of u∗ and P∗ at round t, defined in equation 17.
Absolute constants.
1/(2 + exp(−B) + exp(B)).
Eluder coefficient from Definition 3.
Generalized Eluder-type condition from Definition 4.
Total variation distance between two distributions.

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

A.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) prompt-
ing (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 ma-
nipulations 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 limita-
tions, 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., 2023; Chen et al., 2022). This method
significantly outperforms traditional CoT-based techniques in mathematical problem solving. How-
ever, 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;

18

918
919

920
921
922
923
924
925

926
927
928
929
930
931

932
933
934
935
936

937
938
939
940
941
942

943
944
945
946
947
948

949
950
951
952
953
954

955
956
957
958
959

960
961
962
963
964
965

966
967
968
969
970
971

Under review as a conference paper at ICLR 2025

972
973

974
975
976
977
978
979

980
981
982
983
984
985

986
987
988
989
990

991
992
993
994
995
996

997
998
999
1000
1001
1002

1003
1004
1005
1006
1007
1008

1009
1010
1011
1012
1013

1014
1015
1016
1017
1018
1019

1020
1021
1022
1023
1024
1025

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 (Gemini et al., 2023), and Claude (Anthropic,
2023). However, applying PPO requires extensive efforts and resources (Choshen et al., 2019; En-
gstrom 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., 2024; 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., 2024; 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 prob-
lem 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., 2024; Xu et al.,
2023; Hoang Tran, 2024) and achieve better performance for CoT reasoning. Inspired by the suc-
cess of PRMs, recent studies have explored generating proxy step-wise labels for the intermediate

19

Under review as a conference paper at ICLR 2025

1026
1027

1028
1029
1030
1031
1032
1033

1034
1035
1036
1037
1038
1039

1040
1041
1042
1043
1044

1045
1046
1047
1048
1049
1050

1051
1052
1053
1054
1055
1056

1057
1058
1059
1060
1061
1062

1063
1064
1065
1066
1067

1068
1069
1070
1071
1072
1073

1074
1075
1076
1077
1078
1079

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 frame-
work, 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 2.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 equation 7). 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.

A.3 MISSING DETAILS

Multi-turn KTO. With equation 7 implying that with term (C) = 0, the implicit reward is given
h=1 log π∗
by A = η (cid:80)H
πref,h(ah|sh) , a multi-turn version of KTO (Ethayarajh et al., 2024), denoted as
M-KTO, can also be naturally derived:

h(ah|sh)

LM-KTO(θ) = Ex,y∼D

(cid:2)λy − v(x, y)(cid:3),

(11)

where

and

uθ(x, y) = η

H
(cid:88)

h=1

log

πu,h(ah|sh)
πref,h(ah|sh)

,

z0 = Ex′∼D,τ ′∼πθ(·|x′)

H
(cid:88)

h=1

DKL

(cid:0)πθ(·|sh), πref (·|sh)(cid:1),

v(x, y) =

(cid:40)

λ+σ(cid:0)η(uθ(x, y) − z0)(cid:1)
λ−σ(cid:0)η(z0 − uθ(x, y))(cid:1)

if y ∼ ydesirable|x
if y ∼ yundesirable|x

.

Here λ+ and λ− 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.

B IMPLEMENTATION DETAILS

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 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

20

Under review as a conference paper at ICLR 2025

1080
1081

1082
1083
1084
1085
1086
1087

1088
1089
1090
1091
1092
1093

1094
1095
1096
1097
1098

1099
1100
1101
1102
1103
1104

1105
1106
1107
1108
1109
1110

1111
1112
1113
1114
1115
1116

1117
1118
1119
1120
1121

1122
1123
1124
1125
1126
1127

1128
1129
1130
1131
1132
1133

prompt and the messages returned by the Python interpreter) in the training. 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 format and generation. We format the data into a multi-turn chat where the user asks the
LLMs a question, and provide the messages returned by the Python interpreter in the subsequent
user rounds of chat. 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. 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 oh, 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.

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

Prompt template. We do not tune the prompt though we do observe that the prompt engineering
can further improve the performance. For all the experiments, we simply adopt the chat template of
the models to form a multi-turn chat as in Figure 1.

21

Under review as a conference paper at ICLR 2025

Figure 5: 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.

C ADDITIONAL EXPERIMENTAL RESULTS

We include additional ablation studies in this section for a more comprehensive understanding of
the proposed algorithm.

The best model is obtained with starting checkpoint fine-tuned with more than 1 epochs. Tun-
stall 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 5. 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 5: 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 ↑.

Model

Method

GSM8K MATH

SFT 1 epoch

Gemma-1.1-it-7B
Gemma-1.1-it-7B SFT 1 epoch + Iterative M-DPO 80.6 ↑5.5
Gemma-1.1-it-7B
Gemma-1.1-it-7B SFT 2 epoch + Iterative M-DPO 82.4 ↑7.1
Gemma-1.1-it-7B
Gemma-1.1-it-7B SFT 3 epoch + Iterative M-DPO 83.9 ↑6.4

SFT 2 epoch

SFT 3 epoch

75.3

75.1

77.5

41.1
46.7 ↑5.6
44.0
49.8 ↑5.8
46.1
51.2 ↑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) ca-

22

1134
1135

1136
1137
1138
1139
1140
1141

1142
1143
1144
1145
1146
1147

1148
1149
1150
1151
1152

1153
1154
1155
1156
1157
1158

1159
1160
1161
1162
1163
1164

1165
1166
1167
1168
1169
1170

1171
1172
1173
1174
1175

1176
1177
1178
1179
1180
1181

1182
1183
1184
1185
1186
1187

Under review as a conference paper at ICLR 2025

pabilities 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 6. 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 6: 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 ↑.
Method

GSM8K MATH

Model

Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B

SFT 3 epoch
SFT 3 epoch + Iterative M-DPO
Iterative M-DPO with NLL loss

Gemma-1.1-it-7B
Gemma-1.1-it-7B
Gemma-1.1-it-7B + M-DPO and NLL loss Iteration 1

SFT 100 steps
+ M-DPO Iteration 1

77.5
83.9
81.7 ↓2.2

50.8
57.8
61.0 ↑3.2

46.1
51.2
49.5 ↓1.7

23.7
27.9
30.1 ↑2.2

On-policy sampling and small learning rate mitigate the probability drops in preferred re-
sponses.
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 simi-
lar 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:
1
πθ(yl|x)

∇θLDP O(πθ, πref ) = −η · σ

rθ(x, yl) − rθ(x, yw)

1
πθ(yw|x)

∇θπθ(yw|x) −

(cid:17)(cid:104)

(cid:16)

∇θπθ(yl|x)

(cid:105)
,

where rθ(x, y) = η log πθ(x,y)
π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 πθ(yl|x) << πθ(yw|x) and the optimization becomes unlearning of the rejected
responses. In this case, the probability of the chosen responses cannot increase. This phenomenon

23

1188
1189

1190
1191
1192
1193
1194
1195

1196
1197
1198
1199
1200
1201

1202
1203
1204
1205
1206

1207
1208
1209
1210
1211
1212

1213
1214
1215
1216
1217
1218

1219
1220
1221
1222
1223
1224

1225
1226
1227
1228
1229

1230
1231
1232
1233
1234
1235

1236
1237
1238
1239
1240
1241

Under review as a conference paper at ICLR 2025

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.

D THEORETICAL PROOFS

D.1 THEORETICAL RESULTS

In this following, 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 tar-
get the most challenging scenario with stochastic and unknown transitions, while as aforementioned,
multi-turn mathematical reasoning with external tools falls into a relatively easier regime with de-
terministic transitions. As mentioned in the main paper, we mostly study 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 been 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
π∗ := arg max

J(π) := J(π; M∗, π0),

and introduce the standard notion of regret as

π

Reg(T ) :=

(cid:88)

t∈[T ]

J(π∗) − J(π1

t ),

(12)

(13)

t ]T

which represents the cumulative performance loss over T steps comparing the learned policies
t=1 against the optimal policy π∗. In addition, we consider that a bounded u∗(x, y) ∈ [0, B] for
[π1
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. (2024).
Definition 2 (Policy Improvement Oracle). For any model M = (S, A, H, P, d0, u) and a reference
function πref , we can compute the optimal policy associated with the model [πM,h]H
h=1 iteratively
as in equation 4.

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

ˆut = arg max

ˆu∈U

Lt(ˆu)

:=

(cid:88)

(cid:104)

(cid:105)
z log(σ(ˆu(τ 1) − ˆu(τ 2))) + (1 − z) log(σ(ˆu(τ 2) − ˆu(τ 1)))

,

(14)

(x,τ 1,τ 2,z)∈∪t−1

i=1 Di

and also an MLE of the transition kernel P∗ over a set P as
(cid:88)

ˆPt = arg max

Lt(ˆP) :=

ˆP∈P

(π,τ )∈∪t−1

i=1 Di

log ˆPπ(τ ),

(15)

where Pπ(τ ) denotes the probability of trajectory τ under policy π and transition kernel P. With the
obtained model ˆMt = (ˆut, ˆPt), the Oracle defined in Definition 2 is called with the reference policy
πref set as the initial policy π0, whose output is adopted as the main policy π1
t .

24

1242
1243

1244
1245
1246
1247
1248
1249

1250
1251
1252
1253
1254
1255

1256
1257
1258
1259
1260

1261
1262
1263
1264
1265
1266

1267
1268
1269
1270
1271
1272

1273
1274
1275
1276
1277
1278

1279
1280
1281
1282
1283

1284
1285
1286
1287
1288
1289

1290
1291
1292
1293
1294
1295

Under review as a conference paper at ICLR 2025

Then, we specify how to choose a theoretically sound exploration policy π2
t . The previous work
of Xiong et al. (2024) 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 M,
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 [ ˆQh]H

h=1 and a reference policy πref with the induced policy ˆπ as

ˆπh(ah|sh) ∝ πref,h(ah|sh) · exp

(cid:16) ˆQh(sh, ah)/η

(cid:17)

,

and the corresponding set of V -functions [ ˆVh]H

h=1 as

ˆVh(sh) = Eah∼ˆπh(·|sh)

(cid:104) ˆQh(sh, ah)

(cid:105)

− ηDKL(ˆπh(·|sh), πref,h(·|sh)),

ˆVH+1(sH+1) = 0,

for any comparator policy π, it holds that

J(π) − J(ˆπ)
= Ed0,π,P∗ [u∗(sH , aH )] − Ed0,ˆπ,P∗ [u∗(sH , aH )]

(cid:88)

+

Ed0,π,P∗

(cid:105)
(cid:104) ˆVh+1(sh+1) − ˆQh(sh, ah)

−

(cid:88)

Ed0,ˆπ,P∗

(cid:105)
(cid:104) ˆVh+1(sh+1) − ˆQh(sh, ah)

h∈[H]

(cid:88)

− η ·

h∈[H]

Ed0,π,P∗ [DKL(πh(·|sh), ˆπh(·|sh))] ,

h∈[H]

where the expectation Ed0,π,P∗ is with respect to the prompt and response (i.e., the trajectory) gen-
erated following d0, P∗ and π.

Based on Proposition 1, the exploration policy π2

t is selected as

π2
t = arg max

π

max
(cid:101)u∈ (cid:101)Ut,(cid:101)P∈ (cid:101)Pt

d0,π,(cid:101)P [(cid:101)u(sH , aH )] − E
E
(cid:16)

−

d0,π,(cid:101)P [ˆut(sH , aH )] − E
E

d0,π1

d0,π1

t ,(cid:101)P [(cid:101)u(sH , aH )]

(cid:88)

+

h∈[H]

E

d0,π,(cid:101)P

(cid:104) ˆVt,h+1(sh+1) −

(cid:17)
t ,(cid:101)P [ˆut(sH , aH )]
(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(sh, ah)

(16)

(cid:105)

,

where (cid:101)Ut and (cid:101)Pt are two confidence sets defined as

(cid:101)Ut = {u ∈ U : Lt(u) ≥ Lt(ˆut) − c1 log(|U|T /δ)},
(cid:101)Pt = {P ∈ P : Lt(P) ≥ Lt(ˆPt) − c1 log(|P|T /δ)}
with c1 denoting an absolute constant here. Note that for the theoretical convenience, we have as-
sumed U and P are finite here, which can be extended to the infinite case using standard discretiza-
tion techniques. It can be observed that π2
t 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 P), the uncertainty from the estimation of transitions dimimishes, which leads to a
similar uncertainty measurement adopted in Xiong et al. (2024).

(17)

The following theorem establishes a rigorous guarantee for the regret incurred.
Theorem 2. Assuming u∗ ∈ U and P∗ ∈ P, with probability at least 1 − δ, we have that

Reg(T ) ≲κ−1B(cid:112)dU T log(|U|T /δ) + B2Hξ(dP , T, c2 log(|P|HT /δ))

(cid:88)

(cid:88)

− η ·

Ed0,π∗,P∗

t∈[T ]

h∈[H]

(cid:2)DKL(π∗

h(·|sh), π1

t,h(·|sh))(cid:3) ,

where κ := 1/(2 + exp(−B) + exp(B)), c2 is an absolute constant, dU is the Eluder coefficient
defined in Definition 3 while dP and ξ(·) are from the generalized Eluder-type condition defined in
Definition 4.

25

1296
1297

1298
1299
1300
1301
1302
1303

1304
1305
1306
1307
1308
1309

1310
1311
1312
1313
1314

1315
1316
1317
1318
1319
1320

1321
1322
1323
1324
1325
1326

1327
1328
1329
1330
1331
1332

1333
1334
1335
1336
1337

1338
1339
1340
1341
1342
1343

1344
1345
1346
1347
1348
1349

Under review as a conference paper at ICLR 2025

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 dU is small and the condition is satisfied with
ξ(dP , T, c2 log(|P|HT /δ)) ≲ (cid:112)dP T log(|P|HT /δ), which implies that the regret of theoretical
version of Algorithm 1 is sublinear in T , further evidencing its statistical efficiency.

D.2 PROOF OF PROPOSITION 1

Proof of Proposition 1. For one policy π, starting with V π
M,H+1 = 0, we recursively define its V -
value and Q-value functions on one model M = (S, A, H, P, d0, u) and the reference policy πref
as

Qπ

M,h(sh, ah) :=

(cid:26) u(sH , aH ),

Eoh∼Ph(·|sh,ah)[V π
(cid:2)Qπ

if h = H,
M,h+1(sh+1)],

M,h(sh) := Eah∼πh(·|sh)
V π

M,h(sh, ah) − η · DKL

if h ≤ H − 1,
(cid:0)πh(·|sh), πref,h(·|sh)(cid:1)(cid:3).

It is noted that with the optimal policy πM, QM,h = QπM
M,h. In the following
discussions, we exclusively focus on the model M∗ = (S, A, H, P∗, d0, u∗) with abbreviations
Qπ

M,h and VM,h = V πM

M∗,h and V π

h = V π

h = Qπ

M∗,h.

For any comparator policy π, it holds that

J(π) − J(ˆπ) = Ed0

(cid:104)

(cid:105)
1 (s1) − ˆV1(s1)
V π

− Ed0

(cid:104)
1 (s1) − ˆV1(s1)
V ˆπ

(cid:105)

,

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

Ed0,π1:h−1,P∗

1:h−1

(cid:104)
h (sh) − ˆVh(sh)
V π

(cid:105)

− Ed0,ˆπ1:h−1,P∗

1:h−1

(cid:104)
h (sh) − ˆVh(sh)
V ˆπ

(cid:105)

(a)
= Ed0,π1:h−1,P∗
− Ed0,π1:h−1,P∗
− Ed0,ˆπ1:h−1,P∗
+ Ed0,ˆπ1:h−1,P∗

[Eπh [Qπ
1:h−1
(cid:104)
Eˆπh
(cid:2)Eˆπh
(cid:104)
Eˆπh

1:h−1

1:h−1

(cid:105)

h(sh, ah)] − ηDKL (πh(·|sh), πref,h(·|sh))]
(cid:105)

− ηDKL (ˆπh(·|sh), πref,h(·|sh))
h(sh, ah)(cid:3) − ηDKL(ˆπh(·|sh), πref,h(·|sh))(cid:3)
(cid:105)
− ηDKL(ˆπh(·|sh), πref,h(·|sh))

(cid:105)

(cid:104) ˆQh(sh, ah)
(cid:2)Qˆπ
(cid:104) ˆQh(sh, ah)
(cid:105)
h(sh, ah) − ˆQh(sh, ah)
(cid:104) ˆQh(sh, ah)

Eπh

(cid:105)

− Eˆπh
(cid:123)(cid:122)
term (I)

Qπ
(cid:104)

(cid:124)

1:h−1
(cid:104)

= Ed0,π1:h,P∗

1:h−1

+ Ed0,π1:h−1,P∗

1:h−1

(cid:104)

Qˆπ

h(sh, ah) − ˆQh(sh, ah)

(cid:105)

− Ed0,ˆπ1:h,P∗
(cid:104) ˆQh(sh, ah)

1:h−1
(cid:105)(cid:105)

(cid:125)

1:h−1

− η · Ed0,π1:h−1,P∗
+ η · Ed0,π1:h−1,P∗
1:h−1
(cid:104)
(b)
= Ed0,π1:h,P∗
− η · Ed0,π1:h−1,P∗

1:h−1

Qπ

1:h−1

[DKL (πh(·|sh), πref,h(·|sh))]

[DKL (ˆπh(·|sh), πref,h(·|sh))]

h(sh, ah) − ˆQh(sh, ah)
[DKL (πh(·|sh), ˆπh(·|sh))] .

(cid:105)

− Ed0,ˆπ1:h,P∗

1:h−1

(cid:104)

Qˆπ

h(sh, ah) − ˆQh(sh, ah)

(cid:105)

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

(term I) := Eπh

(cid:105)
(cid:104) ˆQh(sh, ah)

= η · Eπh

(cid:20)

log

− Eˆπh
ˆπh(ah|sh)
πref,h(ah|sh)

(cid:105)
(cid:104) ˆQh(sh, ah)
(cid:20)
(cid:21)

− η · Eˆπh

log

(cid:21)

ˆπh(ah|sh)
πref,h(ah|sh)

= η · DKL (πh(·|sh), πref,h(·|sh)) − η · DKL (πh(·|sh), ˆπh(·|sh))
− η · DKL (ˆπh(·|sh), πref,h(·|sh)) .

26

1350
1351

1352
1353
1354
1355
1356
1357

1358
1359
1360
1361
1362
1363

1364
1365
1366
1367
1368

1369
1370
1371
1372
1373
1374

1375
1376
1377
1378
1379
1380

1381
1382
1383
1384
1385
1386

1387
1388
1389
1390
1391

1392
1393
1394
1395
1396
1397

1398
1399
1400
1401
1402
1403

Under review as a conference paper at ICLR 2025

where the second equation is from the relationship that

ˆQh(sh, ah) = η · log

ˆπh(ah|sh)
πref,h(ah|sh)

− η · log ˆZh(sh).

Furthermore, if h = H, we can obtain that

Ed0,π1:H−1,P∗

1:H−1

1:H−1

= Ed0,π1:H ,P∗
− η · Ed0,π1:H−1,P∗
= Ed0,π1:H ,P∗
+ Ed0,π1:H ,P∗
− η · Ed0,π1:H−1,P∗

1:H−1

1:H

(cid:104)

(cid:105)
H (sH ) − ˆVH (sH )
V π
(cid:104)

− Ed0,ˆπ1:H−1,P∗
(cid:105)
u∗(sH , aH ) − ˆQH (sH , aH )

(cid:104)
H (sH ) − ˆVH (sH )
V ˆπ

(cid:105)

1:H−1

− Ed0,ˆπ1:H ,P∗

1:H−1

(cid:105)
(cid:104)
u∗(sH , aH ) − ˆQH (sH , aH )

[DKL (πH (·|sH ), ˆπH (·|sH ))]

1:H−1

[u∗(sH , aH )] − Ed0,ˆπ1:H ,P∗
1:H−1
(cid:104) ˆVH+1(sH+1) − ˆQH (sH , aH )
(cid:105)

[u∗(sH , aH )]

− Ed0,ˆπ1:H ,P∗

1:H

[DKL (πH (·|sH )||ˆπH (·|sH ))] ,

1:H−1

(cid:104) ˆVH+1(sH+1) − ˆQH (sH , aH )

(cid:105)

where the second equality leverages that ˆVH+1(sH+1) = 0; otherwise, for all h ≤ H − 1, it holds
that

Ed0,π1:h−1,P∗

1:h−1

1:h−1

= Ed0,π1:h,P∗
− η · Ed0,π1:h−1,P∗
= Ed0,π1:h,P∗
− η · Ed0,π1:h−1,P∗
+ Ed0,π1:h,P∗

1:h

1:h

(cid:105)

(cid:104)
h (sh) − ˆVh(sh)
V π
(cid:104)

− Ed0,ˆπ1:h−1,P∗
(cid:105)
h(sh, ah) − ˆQh(sh, ah)

Qπ

[DKL (πh(·|sh)||ˆπh(·|sh))]

1:h−1

(cid:105)
(cid:104)
h (sh) − ˆVh(sh)
V ˆπ

1:h−1

− Ed0,ˆπ1:h,P∗

1:h−1

(cid:104)

Qˆπ

h(sh, ah) − ˆQh(sh, ah)

(cid:105)

(cid:105)
(cid:104) ˆVh+1(sh+1) − ˆQh(sh, ah)

− Ed0,ˆπ1:h,P∗

1:h

(cid:105)
(cid:104) ˆVh+1(sh+1) − ˆQh(sh, ah)

[DKL (πh(·|sh)||ˆπh(·|sh))]

1:h−1

(cid:104)
h+1(sh+1) − ˆVh+1(sh+1)
V π

(cid:105)

− Ed0,π1:h,P∗

1:h

(cid:104)

(cid:105)
h+1(sh+1) − ˆVh+1(sh+1)
V ˆπ

.

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

D.3 PROOF OF THEOREM 2

First, with the assumption u∗ ∈ U and P∗ ∈ P, the following lemma demonstrates that (cid:101)Ut and (cid:101)Pt
are valid confidence sets.
Lemma 1 (Proposition B.1 from Liu et al. (2023a)). There exists an absolute constant c1 such that
for any δ ∈ (0, 1], with probability at least 1 − δ, for all t ∈ [T ], ˆu ∈ U, and ˆP ∈ P, it holds that

Lt(ˆu) − Lt(u∗) ≤ c1 log(|U|T /δ),

Lt(ˆP) − Lt(P∗) ≤ c1 log(|P|T /δ),

which implies that u∗ ∈ (cid:101)Ut and P∗ ∈ (cid:101)Pt.

Then, we provide an additional lemma demonstrating the in-sample error of the MLE and optimistic
estimators.
Lemma 2. There exists an absolute constant c2 such that for any δ ∈ (0, 1], with probability at least
1 − δ, for all t ∈ [T ], we have

(cid:12)
(cid:12)σ (cid:0)ˆut(s2

i,H , a2

i,H ) − ˆut(s1

i,H , a1

i,H )(cid:1) − σ (cid:0)u∗(s2

i,H , a2

i,H ) − u∗(s1

i,H , a1

i,H )(cid:1)(cid:12)
(cid:12)

(cid:12)
(cid:12)σ (cid:0)

(cid:101)ut(s2

i,H , a2

i,H ) − (cid:101)ut(s1

i,H , a1

i,H )(cid:1) − σ (cid:0)u∗(s2

i,H , a2

i,H ) − u∗(s1

i,H , a1

i,H )(cid:1)(cid:12)
(cid:12)

2

2

≤ c2 log(|U|T /δ);

≤ c2 log(|U|T /δ),

(cid:88)

i<t
(cid:88)

i<t

and for all t ∈ [T ], h ∈ [H], we have

(cid:88)

(cid:88)

(cid:88)

(cid:16)

TV

j∈{1,2}

h∈[H]

i<t

{d0, πj

i , [P∗

1:h−1, ˆPt,h, P∗

h+1:H ]}, {d0, πj

i , P∗

1:H }

(cid:17)2

≤ c2 log(|P|HT /δ);

27

1404
1405
1406
1407

1408
1409
1410
1411
1412
1413

1414
1415
1416
1417
1418

1419
1420
1421
1422
1423
1424

1425
1426
1427
1428
1429
1430

1431
1432
1433
1434
1435
1436

1437
1438
1439
1440
1441

1442
1443
1444
1445
1446
1447

1448
1449
1450
1451
1452
1453

1454
1455
1456
1457

Under review as a conference paper at ICLR 2025

(cid:88)

(cid:88)

(cid:88)

(cid:16)

TV

j∈{1,2}

h∈[H]

i<t

{d0, πj

i , [P∗

1:h−1, (cid:101)Pt,h, P∗

h+1:H ]}, {d0, πj

i , P∗

1:H }

(cid:17)2

≤ c2 log(|P|HT /δ),

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

Proof of Lemma 2. First, for (cid:101)ut, we can obtain that with probability at least 1 − δ, there exists an
absolute constant c such that for all t ∈ [T ],

(cid:12)
(cid:12)σ (cid:0)

(cid:88)

i<t

(cid:101)ut(s2

i,H , a2

i,H ) − (cid:101)ut(s1

i,H , a1

i,H )(cid:1) − σ (cid:0)u∗(s2

i,H , a2

i,H ) − u∗(s1

i,H , a1

i,H )(cid:1)(cid:12)
2
(cid:12)

≤ c ·

(cid:88)

i<t

log

zi · σ (cid:0)u∗(s1
zi · σ (cid:0)
(cid:101)ut(s1

i,H , a1
i,H , a1

i,H ) − u∗(s2
i,H ) − (cid:101)ut(s2

i,H , a2
i,H , a2

i,H )(cid:1) + (1 − zi) · σ (cid:0)u∗(s2
i,H )(cid:1) + (1 − zi) · σ (cid:0)
(cid:101)ut(s2

i,H , a2
i,H , a2

i,H ) − u∗(s1
i,H ) − (cid:101)ut(s1

i,H , a1
i,H , a1

i,H )(cid:1)
i,H )(cid:1)

+ c · log(|U|T /δ)
= c (Lt(u∗) − Lt((cid:101)ut) + log(|U|T /δ))
≤ c (Lt(u∗) − Lt(ˆut) + c1 log(|U|T /δ) + log(|U|T /δ))
≤ c2 log(|U|T /δ).

where the first inequality is from Proposition B.2 from Liu et al. (2023a) and the second inequality
uses Lemma 1. The result for ˆut can be similarly established.
Then, following similar steps, for (cid:101)Pt, we can obtain that with probability at least 1 − δ, there exists
an absolute constant c such that for all t ∈ [T ],

(cid:88)

(cid:88)

(cid:88)

(cid:16)

TV

{d0, πj

i , [P∗

1:h−1, (cid:101)Pt,h, P∗

h+1:H ]}, {d0, πj

i , P∗

1:H }

(cid:17)2

j∈{1,2}

h∈[H]

i<t

(cid:88)

(cid:88)

≤

(cid:32)

(cid:88)

c ·

log

h∈[H]

i<t

j∈{1,2}


h(sj
P∗
(cid:101)Pt,h(sj

i,h+1|sj
i,h+1|sj

i,h, aj
i,h, aj

i,h)
i,h)

(cid:33)

+ log(|Ph|HT /δ)



= c ·



(cid:16)

= c ·

log

(cid:88)

(cid:88)

P∗,πj
i (τ j
i )
(cid:101)Pπj
t (τ j
i )
Lt(P∗) − Lt((cid:101)Pt) + 2 log(|P|HT /δ)

j∈{1,2}

i<t

i

(cid:17)

+ 2 log(|P|HT /δ)



≤ c ·

(cid:16)

(cid:17)
Lt(P∗) − Lt(ˆPt) + c1 log(|P|T /δ) + 2 log(|P|HT /δ)

≤ c2 log(|P|HT /δ).

The result for ˆPt can also be similarly established.

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

J(π∗) − J(π1
t )
= Ed0,π∗,P∗ [u∗(sH , aH )] − E

(cid:88)

+

Ed0,π∗,P∗

h∈[H]

d0,π1
(cid:104) ˆVt,h+1(sh+1) −

t ,P∗ [u∗(sH , aH )] −
(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(sh, ah)

Ed0,π∗,P∗ [ˆut(sH , aH )] − E
(cid:105)

(cid:88)

d0,π1
(cid:104) ˆVt,h+1(sh+1) −

t ,P∗ [ˆut(sH , aH )]
(cid:104)ˆPt,h ˆVt,h+1

−

E

(cid:105)

(cid:105)

(sh, ah)

d0,π1

t ,P∗

(cid:16)

(cid:17)

d0,π2

t ,(cid:101)Pt

[(cid:101)ut(sH , aH )] − E

d0,π1

t ,(cid:101)Pt

≤ E
(cid:124)

(cid:16)

E

[(cid:101)ut(sH , aH )] −
(cid:123)(cid:122)
term (I)t

h∈[H]

d0,π2

t ,(cid:101)Pt

[ˆut(sH , aH )] − E

d0,π1

t ,(cid:101)Pt

[ˆut(sH , aH )]

(cid:17)

(cid:125)

E

d0,π2

t ,(cid:101)Pt

(cid:88)

+

h∈[H]
(cid:124)

(cid:104) ˆVt,h+1(sh+1) −

(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(sh, ah)

(cid:105)

+

(cid:88)

E

(cid:104)(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah) − ˆVt,h+1(sh+1)

,

d0,π1

t ,P∗

h∈[H]

(cid:123)(cid:122)
term (II)t

(cid:125)

where the inequality is from the definition of π2

t and the fact that (u∗, P∗) ∈ (cid:101)Ut × (cid:101)Pt from Lemma 1.

28

1458
1459
1460
1461

1462
1463
1464
1465
1466
1467

1468
1469
1470
1471
1472

1473
1474
1475
1476
1477
1478

1479
1480
1481
1482
1483
1484

1485
1486
1487
1488
1489
1490

1491
1492
1493
1494
1495

1496
1497
1498
1499
1500
1501

1502
1503
1504
1505
1506
1507

1508
1509
1510
1511

Under review as a conference paper at ICLR 2025

We define the following terms:

term (A)t := E

d0,π2

term (B)t := E

t ,P∗ [(cid:101)ut(sH , aH )] − E
t ,P∗ [u∗(sH , aH )] − E

term (C)t :=

term (D)t :=

d0,π2
(cid:88)

(cid:88)

j∈{1,2}
(cid:88)

h∈[H]
(cid:88)

j∈{1,2}

h∈[H]

(cid:104)

(cid:104)

E

E

d0,πj

t ,P∗

d0,πj

t ,P∗

d0,π1

t ,P∗ [(cid:101)ut(sH , aH )] −
t ,P∗ [u∗(sH , aH )] −
(cid:101)Pt,h(·|sh, ah), P∗

d0,π1
(cid:16)

TV

(cid:16)

E

(cid:16)

E

d0,π2

d0,π2
(cid:17)(cid:105)

h(·|sh, ah)

,

(cid:16)ˆPt,h(·|sh, ah), P∗

h(·|sh, ah)

(cid:17)(cid:105)

.

TV

t ,P∗ [u∗(sH , aH )] − E
t ,P∗ [ˆut(sH , aH )] − E

d0,π1

t ,P∗ [u∗(sH , aH )]

d0,π1

t ,P∗ [ˆut(sH , aH )]

(cid:17)

(cid:17)

,

,

For term (I)t, we have that

term (I)t := E

d0,π2

[(cid:101)ut(sH , aH )] −

(cid:16)

E

d0,π2

d0,π1

t ,(cid:101)Pt

= E

d0,π2

+ E

+ E

d0,π2

d0,π2

+ E

d0,π2

≤ E

d0,π2

+ E

d0,π2

[(cid:101)ut(sH , aH )] − E
t ,(cid:101)Pt
t ,P∗ [(cid:101)ut(sH , aH )] − E
t ,P∗ [u∗
t (sH , aH )] − E

[(cid:101)ut(sH , aH )] − E
t ,(cid:101)Pt
t ,P∗ [ˆut(sH , aH )] − E
t ,P∗ [(cid:101)ut(sH , aH )] − E
t ,P∗ [u∗
t (sH , aH )] − E
(cid:16)

+ 4B · TV

{d0, π1

t , (cid:101)Pt}, {d0, π1

d0,π1

d0,π1

t ,P∗ [(cid:101)ut(sH , aH )] −
t ,P∗ [u∗

t (sH , aH )] −
(cid:16)

[(cid:101)ut(sH , aH )] −

d0,π1

t ,(cid:101)Pt

d0,π1

t ,P∗ [ˆut(sH , aH )] −

d0,π1

d0,π1

t ,P∗ [(cid:101)ut(sH , aH )] −
t ,P∗ [u∗
(cid:17)
t , P∗}

t (sH , aH )] −
(cid:16)

+ 4B · TV

(cid:16)

E

(cid:16)

E

d0,π2

d0,π2

E

d0,π2

(cid:16)

E

(cid:16)

E

(cid:16)

E

d0,π2

d0,π2

d0,π2

[ˆut(sH , aH )] − E

t (sH , aH )] − E

t ,(cid:101)Pt
t ,P∗ [u∗
t ,P∗ [ˆut(sH , aH )] − E
t ,P∗ [(cid:101)ut(sH , aH )] − E
[ˆut(sH , aH )] − E

t (sH , aH )] − E

t ,(cid:101)Pt
t ,P∗ [u∗
t ,P∗ [ˆut(sH , aH )] − E
t , (cid:101)Pt}, {d0, π2

t , P}

(cid:17)

{d0, π2
(cid:16)

d0,π2

t ,P∗ [(cid:101)ut(sH , aH )] − E

d0,π1

t ,P∗ [(cid:101)ut(sH , aH )] −

E

d0,π2

t ,P∗ [u∗

t (sH , aH )] − E

d0,π2

t ,P∗ [u∗

t (sH , aH )] − E

d0,π1

t ,P∗ [u∗

t (sH , aH )] −

d0,π2

t ,P∗ [ˆut(sH , aH )] − E

(cid:123)(cid:122)
term (A)t
(cid:16)

E

≤ E
(cid:124)

+ E
(cid:124)

(cid:17)

(cid:17)

[ˆut(sH , aH )]

d0,π1

t ,(cid:101)Pt
t ,P∗ [u∗

d0,π1

d0,π1

t (sH , aH )]
(cid:17)
t ,P∗ [ˆut(sH , aH )]
(cid:17)
t ,P∗ [(cid:101)ut(sH , aH )]
(cid:17)

[ˆut(sH , aH )]

d0,π1

d0,π1

t ,(cid:101)Pt
t ,P∗ [u∗

d0,π1

(cid:17)

t (sH , aH )]
(cid:17)

d0,π1

t ,P∗ [ˆut(sH , aH )]

d0,π1

t ,P∗ [u∗

t (sH , aH )]

(cid:17)

(cid:125)

d0,π1

t ,P∗ [ˆut(sH , aH )]

(cid:17)

(cid:125)

Ed0

E

πj
t ,P∗

(cid:104)
TV

(cid:16)

(cid:101)Pt,h(·|sh, ah), P∗

h(·|sh, ah)

(cid:123)(cid:122)
term (B)t

(cid:123)(cid:122)
term (C)t

(cid:17)(cid:105)

.

(cid:125)

+ 4B ·

(cid:88)

(cid:88)

h∈[H]

j∈{1,2}
(cid:124)

For term (II)t, we have that
(cid:88)

term (II)t =

E

d0,π2

t ,(cid:101)Pt

E

d0,π1

t ,P∗

E

d0,π2

t ,P∗

E

d0,π2

t ,(cid:101)Pt

E

d0,π2

t ,P∗

E

d0,π1

t ,P∗

(cid:104) ˆVt,h+1(sh+1) −

(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah)

(cid:104)(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah) − ˆVt,h+1(sh+1)

(cid:104) ˆVt,h+1(sh+1) −

(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah)

(cid:104) ˆVt,h+1(sh+1) −

(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah)

(cid:104) ˆVt,h+1(sh+1) −

(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah)

(cid:104)(cid:104)ˆPt,h ˆVt,h+1

(cid:105)

(cid:105)
(sh, ah) − ˆVt,h+1(sh+1)

(cid:88)

(cid:88)

E

d0,πj

t ,P∗

(cid:104)
TV(ˆPt,h(·|sh, ah)), P∗

(cid:105)
h(·|sh, ah)

h∈[H]
(cid:88)

h∈[H]
(cid:88)

h∈[H]
(cid:88)

h∈[H]
(cid:88)

h∈[H]
(cid:88)

h∈[H]

+

=

+

−

+

≤ 2B ·

j∈{1,2}

h∈[H]

29

1512
1513

1514
1515
1516
1517
1518
1519

1520
1521
1522
1523
1524
1525

1526
1527
1528
1529
1530

1531
1532
1533
1534
1535
1536

1537
1538
1539
1540
1541
1542

1543
1544
1545
1546
1547
1548

1549
1550
1551
1552
1553

1554
1555
1556
1557
1558
1559

1560
1561
1562
1563
1564
1565

Under review as a conference paper at ICLR 2025

t , (cid:101)Pt}, {d0, π2
+ 2BH · TV({d0, π2
(cid:88)
E

(cid:88)

≤ 2B ·

d0,πj

t ,P∗

t , P∗})

(cid:104)
TV(ˆPt,h(·|sh, ah)), P∗

(cid:105)
h(·|sh, ah)

h∈[H]

j∈{1,2}
(cid:124)

+ 2BH ·

(cid:88)

(cid:88)

h∈[H]

j∈{1,2}
(cid:124)

E

d0,πj

t ,P∗

(cid:123)(cid:122)
term (D)t
(cid:104)
TV((cid:101)Pt,h(·|sh, ah)), P∗

h(·|sh, ah)

(cid:125)

(cid:105)

.

(cid:125)

(cid:123)(cid:122)
term (C)t

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

TV({d0, π2

t , (cid:101)Pt}, {d0, π2

t , P∗}) ≤

(cid:88)

h∈[H]

E

d0,π2

t ,P∗

(cid:104)

TV

(cid:16)

(cid:101)Pt,h(·|sh, ah), P∗

h(·|sh, ah)

(cid:17)(cid:105)

,

which can be derived as
TV({d0, π2

t , (cid:101)Pt}, {d0, π2

t , P∗}) ≤

(cid:88)

TV

(cid:16)
{d0, π2

t , P∗

1:h−1, (cid:101)Pt,h:H }, {d0, π2

t , P∗

1:h, (cid:101)Pt,h+1:H }

(cid:17)

h∈[H]
(cid:88)

h∈[H]

=

E

d0,π2

t ,P∗

(cid:104)
TV

(cid:16)

(cid:101)Pt,h(·|sh, ah), P∗

h(·|sh, ah)}

(cid:17)(cid:105)

.

Then, we can obtain that

(cid:88)

t∈[T ]

J(π∗) − J(ˆπ1

t ) ≤

(cid:88)

t∈[T ]

term (A)t +

(cid:88)

t∈[T ]

term (B)t

+ (4B + 2BH)

(cid:88)

t∈[T ]

term (C)t + 2B

(cid:88)

t∈[T ]

term (D)t.

Then, we control the sum of each individual term in the following. First, for term (A)t, with proba-
bility at least 1 − δ, we have that
(cid:88)

term (A)t

t∈[T ]

=

≤

≤

(cid:88)

t∈[T ]
(cid:88)

E

d0,π2

t ,P∗ [(cid:101)ut(sH , aH )] − E

d0,π1

t ,P∗ [(cid:101)ut(sH , aH )] −

(cid:16)

E

d0,π2

t ,P∗ [u∗(sH , aH )] − E

d0,π1

(cid:17)
t ,P∗ [u∗(sH , aH )]

(cid:101)ut(s2

t,H , a2

t,H ) − (cid:101)ut(s1

t,H , a1

t,H ) − (cid:0)u∗(s2

t,H , a2

t,H ) − u∗(s1

t,H , a1

t,H )(cid:1) + O(B(cid:112)T log(1/δ))

t∈[T ]
(cid:118)
(cid:117)
(cid:117)
(cid:116)dU

(cid:32)

1 +

T
(cid:88)

t=2

t−1
(cid:88)

i=1

(cid:0)
(cid:101)ut(s2

i,H , a2

i,H ) − (cid:101)ut(s1

i,H , a1

i,H ) − (cid:0)u∗(s2

i,H , a2

i,H ) − u∗(s1

i,H , a1

i,H )(cid:1)(cid:1)2

(cid:33)

+ O(B(cid:112)T log(1/δ))

(cid:118)
(cid:117)
(cid:117)
(cid:116)dU

≤

T
(cid:88)

t=2

(cid:32)

1 + κ−2

t−1
(cid:88)

i=1

(cid:0)σ (cid:0)

(cid:101)ut(s2

i,H , a2

i,H ) − (cid:101)ut(s1

i,H , a1

i,H )(cid:1) − σ (cid:0)u∗(s2

i,H , a2

i,H ) − u∗(s1

i,H , a1

i,H )(cid:1)(cid:1)2

(cid:33)

+ O(B(cid:112)T log(1/δ))
≲ κ−1B(cid:112)dU T log(|U|T /δ),
where the first inequality is from the Hoeffding inequality, the second inequality uses the Eluder
coefficient dU := EC(1, U − U, T ) from Definition 3, the third inequality leverages the mean value
theorem with κ := 1/(2 + exp(−B) + exp(B)) representing the minimum derivative of σ(·) in the
regime of [0, B], and the last inequality incorporates Lemma 2. A similar result can be obtained for
term (B)t.
For term (C)t, we have that
(cid:88)

(cid:88)

(cid:88)

(cid:88)

(cid:17)(cid:105)

(cid:16)

(cid:104)

E

d0,πj

t ,P∗

TV

(cid:101)Pt,h(·|sh, ah), P∗

h(·|sh, ah)

term (C)t =

t∈[T ]

j∈{1,2}

t∈[T ]

h∈[H]

30

1566
1567

1568
1569
1570
1571
1572
1573

1574
1575
1576
1577
1578
1579

1580
1581
1582
1583
1584

1585
1586
1587
1588
1589
1590

1591
1592
1593
1594
1595
1596

1597
1598
1599
1600
1601
1602

1603
1604
1605
1606
1607

1608
1609
1610
1611
1612
1613

1614
1615
1616
1617
1618
1619

Under review as a conference paper at ICLR 2025

(cid:88)

(cid:88)

(cid:88)

=

(cid:16)

TV

{d0, πj

t , [P∗

1:h−1, (cid:101)Pt,h, P∗

h+1:H ]}, {d0, πj

t , P∗

1:H }

(cid:17)

t∈[T ]

j∈{1,2}

h∈[H]
≤ 2H · ξ(dP , T, c2 log(|P|HT /δ)),

where the last step is from the generalized Eluder-type condition in Definition 4 and Lemma 2. A
similar result can be obtained for term (D)t.
Finally, we obtain that

Reg(T ) ≲κ−1B(cid:112)dU T log(|U|T /δ) + B2Hξ(dP , T, c2 log(|P|HT /δ)

(cid:88)

(cid:88)

− η ·

Ed0,π∗,P∗

t∈[T ]

h∈[H]

which concludes the proof.

E TECHNICAL LEMMAS

(cid:2)DKL(π∗

h(·|sh), π1

t,h(·|sh))(cid:3) ,

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(·|x), written as

(cid:104)

Ew∼p(·)

− U (w) + ηDKL

= ηDKL

(cid:16)

p(·), p0(·) exp

(cid:0)p(·), p0(·)(cid:1)(cid:105)
(cid:17)(cid:17)
(cid:16) 1
η

U (·)

− η · log Ew∼p0(·) exp (cid:0) 1
η

(cid:124)

(cid:123)(cid:122)
Cr

U (w)(cid:1)

,

(cid:125)

where the minimizer of the loss functional is p∗(w) = 1
Cr
distribution.
Definition 3 (Eluder Coefficient, Definition 17.17 in Zhang (2023)). Given a function class F,
its Eluder coefficient EC(λ, F, T ) is defined as the smallest number d so that for any sequence
{xt : t ∈ [T ]} and {ft : t ∈ [T ]} ∈ F,

, also known as Gibbs

p0(w) exp

(cid:16) 1

(cid:17)
η U (w)

T
(cid:88)

t=2

|ft(xt) − f ∗(xt)| ≤

(cid:118)
(cid:117)
(cid:117)
(cid:116)d

T
(cid:88)

t=2

(cid:32)

λ +

t−1
(cid:88)

(ft(xi) − f ∗(xi))2

(cid:33)
.

i=1

Definition 4 (Generalized Eluder-type Condition, Condition 3.1 in Liu et al. (2023a)). There exists a
real number dP ∈ R+ and a function ξ such that for any (T, ∆) ∈ N×R+, transitions {P′
t : t ∈ [T ]}
and policies {πt : t ∈ [T ]}, we have

∀t ∈ [T ],

(cid:88)

i<t

⇒

TV({d0, P′

i, πi}, {d0, P, πi})2 ≤ ∆

(cid:88)

TV({d0, P′

t, πt}, {d0, P, πt}) ≤ ξ(dP , T, ∆).

t∈[T ]

31

1620
1621

1622
1623
1624
1625
1626
1627

1628
1629
1630
1631
1632
1633

1634
1635
1636
1637
1638

1639
1640
1641
1642
1643
1644

1645
1646
1647
1648
1649
1650

1651
1652
1653
1654
1655
1656

1657
1658
1659
1660
1661

1662
1663
1664
1665
1666
1667

1668
1669
1670
1671
1672
1673

