Q-SFT: Q-Learning for Language Models via Supervised Fine-Tuning

## Abstract

Abstract Value-based reinforcement learning (RL) can in principle learn effective policies for a wide range of multi-turn problems, from games to dialogue to robotic control, including via offline RL from static previously collected datasets. However, despite the widespread use of policy gradient methods to train large language models for single turn tasks (e.g., question answering), value-based methods for multi-turn RL in an off-policy or offline setting have proven particularly challenging to scale to the setting of large language models. This setting requires effectively leveraging pretraining, scaling to large architectures with billions of parameters, and training on large datasets, all of which represent major challenges for current value-based RL methods.
In this work, we propose a novel offline RL algorithm that addresses these drawbacks, casting Q-learning as a modified supervised fine-tuning (SFT) problem where the probabilities of tokens directly translate to Q-values.
In this way we obtain an algorithm that smoothly transitions from maximizing the likelihood of the data during pretraining to learning a near-optimal Q-function during finetuning.
Our algorithm has strong theoretical foundations, enjoying performance bounds similar to state-of-the-art Q-learning methods, while in practice utilizing an objective that closely resembles SFT.
Because of this, our approach can enjoy the full benefits of the pretraining of language models, without the need to reinitialize any weights before RL finetuning, and without the need to initialize new heads for predicting values or advantages.
Empirically, we evaluate our method on both pretrained LLMs and VLMs, on a variety of tasks including both natural language dialogue and robotic manipulation and navigation from images.

## 1 Introduction

Recently, some of the most impressive feats in AI have been performed through language models,
which are pretrained on large-scale data and adapted to a wide range of downstream tasks  (Bommasani et al., ).
Many of these tasks, such as natural language dialogue or robotic control, require complex sequential decision-making.
Reinforcement learning (RL) Sutton & Barto ()
is a powerful paradigm for solving such tasks (Mnih et al., ; Silver et al., ; AlphaStar, ).
Furthermore, offline RL Levine et al. () has been shown to do so from only static datasets, such as suboptimal demonstrations from any unknown behavior policy, without the need for any additional interaction.
Though offline RL has been used to fine-tune large language models (LLMs) or vision language models (VLMs) (Ouyang et al., ; Bai et al., ), its usefulness has been limited to generating better single responses rather than multi-turn, sequential scenarios where RL should theoretically shine.
For example, across various dialogue tasks, offline RL fine-tuning of LLMs does not reliably outperform supervised fine-tuning (SFT) (Sodhi et al., ; Abdulhai et al., ).
Furthermore, in the realm of navigation and control, popular VLMs are still fine-tuned for multi-task control using SFT (Brohan et al., ; ; Collaboration et al., ).

Single-step problems, such as answering questions, can be tackled with policy gradient methods  (Ouyang et al., ; Rafailov et al., ), but sequential problems, such as dialogue or robotic control, require sample-efficient methods that can utilize data to reason about the dynamics of the problem, which typically requires training value functions (Abdulhai et al., ; Hong et al., ).
Although there are many effective value-based RL methods that could be applied to LLMs and VLMs, in practice such methods have been difficult to adapt to these models with the same effectiveness as policy gradients. We posit that this is due in part to a mismatch between the pretraining objective that these models use, i.e. maximum likelihood estimation, and the fine-tuning objective necessary to train value functions.
This discrepancy means that fine-tuning using multi-turn RL may require discarding some of knowledge gained by maximum likelihood pretraining of LLMs and VLMs, including a broad understanding of language, vision, and even sequential reasoning.

Specifically, we hypothesize two reasons for why fine-tuning foundation models using offline RL is unsuitable in practice.
First, typical offline RL methods require regressing value functions that estimate how appropriate actions, such as an utterance in dialogue, are.
Such algorithms, known as Q-learning,
have achieved impressive results when applied on small networks  (AlphaStar, ; Mnih et al., ), but surprisingly attain disappointing performance when scaled to larger ones  (Sodhi et al., ).
Recent work has attributed this lack of scaling to instability in the value-learning objective, namely in regression towards non-stationary values (Farebrother et al., ).
More importantly, a major advantage of SFT is the potential to leverage existing capabilities of large pretrained models to drastically improve the efficiency when learning a new downstream task.
However, language models are trained to predict likelihoods, but Q-learning instead aims to predict action values; therefore, when fine-tuning, Q-learning algorithms discard the learned likelihoods in favor of only utilizing the underlying representations, which eliminates some of the useful prior knowledge within the pretrained models.
We illustrate this in Figure 1, where value functions are trained must be trained via a new head with reset weights.

In this work, we propose a new algorithm that remedies both drawbacks. Our key insight is simple: *by adding weights to the traditional supervised fine-tuning objective, we can learn probabilities
that conservatively estimate the value function instead of the behavior policy*.
In practice, our approach is implemented by adding weights to the maximum likelihood objective, yielding a *weighted cross entropy loss* where weights are target action values computed from the Bellman recurrence relations.
By using this objective, we are able to avoid the unstable regression objective commonly used in value learning, as well as directly leverage the initial likelihoods resulting from large-scale pretraining.
Theoretically, we can show that such objective results in learned likelihoods that are a product of the data distribution and Q-values, and that our approach is principled and results in performance bounds competitive with other state-of-the-art approaches.
Empirically, we demonstrate the effectiveness of our method on a variety of tasks involving both LLMs, such as language games and dialogue, as well as VLMs, such as navigation and robotic manipulation.

Figure: Figure 1: Our proposed approach allows us to directly leverage the logits from a pretrained model to train value functions. Prior approaches require separately initializing a value head.
Refer to caption: x1.png

## 2 Related Work

Much of the recent work on reinforcement learning (RL) finetuning of LLMs and VLMs uses policy gradient methods and reward models learned from human feedback (e.g., RLHF) (Ziegler et al., ; Stiennon et al., ; Wu et al., ; Nakano et al., ; Bai et al., ; Christiano et al., ; Rafailov et al., ), or from handcrafted AI systems (e.g., RLAIF) (Bai et al., ), to generate better responses to various queries.
However, there is a large discrepancy in the capabilities required to perform self-contained responses in single-step tasks, such as question-answering, and responses in a multi-turn scenarios, such as dialogue. Namely, the latter requires planning to optimize a long-term objective,.
Various prior works provide evidence that existing fine-tuning methods are insufficient to enable language models with such planning capabilities  (Bachmann & Nagarajan, ).

In principle, value-based RL (Lange et al., ; Levine et al., ), specifically Q-learning, can learn effective policies for multi-step tasks that outperform pure imitation via supervised fine-tuning (Kumar et al., ).
Many offline RL algorithms exist that reap the benefits of value-based RL using only static datasets, such as those currently used to fine-tune language models.
Though offline RL algorithms require handling distribution shift (Kumar et al., ),
where the learned policy selects out-of-distribution (OOD) actions with unpredictable consequences, many methods exist that effectively tackle this challenge (Kumar et al., ; Kostrikov et al., ; Kidambi et al., ; Yu et al., ; ).
Due to the promising benefits of offline RL on learning from demonstrations, algorithms have been proposed for learning LLM policies to some success in robotic manipulation (Chebotar et al., ) and language tasks (Snell et al., ).
However, recent evaluation has shown that, on a variety of natural language tasks, Q-learning approaches are often outperformed by supervised ones (Sodhi et al., ; Abdulhai et al., ).
We hypothesize this is due to the mismatch between value-based RL fine-tuning and maximum likelihood pretraining, and propose a new approach that remedies this core issue.

There also exist a paradigm of supervised approaches called return conditioned supervised learning (RCSL), which learn conditional policies on return via a supervised learning objective  (Brandfonbrener et al., ).
The most notable algorithm is Decision Transformer (DT) (Chen et al., ), which can train LLM policies that outperform traditional offline RL methods that rely on Q-learning.
Though it performs well in practice, there is theoretical evidence that the ceiling of performance of such algorithms is below that of value-based offline RL. Specifically, Brandfonbrener et al. () showed that DT and similar approaches can only identify the optimal policy under stronger conditions on the offline data than value-based RL. Our proposed algorithm is similar to RCSL in that we also use a maximum likelihood loss, but we learn values and reap the theoretical benefits of other value-based methods.

Recently, prior attempts have also been made to improve value-based RL algorithms for fine-tuning language models.
Chebotar et al. () propose Q-learning with transformer value functions in manipulation and control tasks by converting actions to sequences of tokens.
We adopt their insight when evaluating on robotics tasks, but use a fundamentally different objective to learn values.
Most similar to ours, Farebrother et al. () propose to replace the regression loss from Q-learning with a cross-entropy loss by casting value learning as a classification problem.
However, while the proposed method also converts value functions to distributions, these likelihoods are not naturally derived from the logits obtained from large-scale pretraining, and must instead be learned from scratch via a separate head with reset weights.
Therefore, like traditional Q-learning, they also suffer from being unable to leverage pretraining efficiently, unlike our approach whose likelihoods are directly initialized by the logits of pretrained LLMs or VLMs.

## 3 Preliminaries

Our work proposes a new RL algorithm for fine-tuning language models, specifically for multi-turn tasks such as dialogue or manipulation and control. Language models operate over a discrete vocabulary of tokens $\mathcal{V}$, and are trained to maximize the likelihood the best next-token $x_{m+1}$ given an input sequence $(x_{0},\ldots,x_{m})$ of tokens, given by $\pi(x_{m+1}|x_{0},\ldots,x_{m})$.
In a multi-turn task such a dialogue, the tokens are words that are chained to form utterances, and the best next-token requires complex, sequential reasoning to understand the utterances so far and plan for the next one. Traditionally, this kind of reasoning can be learned via reinforcement learning (RL).

#### RL fundamentals.

RL aims to optimize agents that interact with a Markov Decision Process (MDP) defined by a tuple $({\mathcal{S}},{\mathcal{A}},P,r,\mu_{1},\gamma)$, where ${\mathcal{S}}$ represents the set of all possible states, ${\mathcal{A}}$ is the set of possible actions, $\mu_{1}$ is the initial state distribution, and $\gamma$ is the discount factor. When action $a\in{\mathcal{A}}$ is executed at state $s\in{\mathcal{S}}$, the next state is generated according to $s^{\prime}\sim P(\cdot|s,a)$, and the agent receives stochastic reward with mean $r(s,a)\in[0,1]$.

The Q-function $Q^{\pi}(s,a)$ for a policy $\pi(\cdot|s)$ represents the discounted long-term reward attained by executing $a$ given observation history $s$ and then following policy $\pi$ thereafter.
$Q^{\pi}$ satisfies the Bellman recurrence:

$$ $Q^{\pi}(s,a)=r(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim P(\cdot|s,a),a^{\prime} \sim\pi(\cdot|s^{\prime})}\left[Q^{\pi}(s^{\prime},a^{\prime})\right]\,.$ $$

The value function $V^{\pi}$ considers the expectation of the Q-function over the policy $V^{\pi}(h)=\mathbb{E}_{a\sim\pi(\cdot|s)}\left[Q^{\pi}(s,a)\right]$.
Meanwhile, the Q-function of the optimal policy $Q^{*}$ satisfies:

$$ $Q^{*}(s,a)=r(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim P(\cdot|s,a)}\left[\max_{a^ {\prime}}Q^{*}(s^{\prime},a^{\prime})\right]\,,$ $$

and the optimal value function is $V^{*}(s)=\max_{a}Q^{*}(s,a)$.
Finally, the expected cumulative reward is given by $J(\pi)=\mathbb{E}_{s_{1}\sim\mu_{1}}\left[V^{\pi}(s_{1})\right]$.
The goal of RL is to optimize a policy $\pi(\cdot\mid s)$ that maximizes the cumulative reward
$J(\pi)=\mathbb{E}_{\mu_{1}}\left[V^{\pi}(s_{1})\right]$.

In offline RL, we are provided with a dataset $\mathcal{D}=\{(s_{i},a_{i},r_{i},s^{\prime}_{i})\}_{i=1}^{N}$ of size $|\mathcal{D}|=N$. We assume that the dataset $\mathcal{D}$ is generated i.i.d. from an effective behavior policy $\pi_{\beta}(a|s)$.
Many state-of-the-art offline RL methods build on Q-learning, which trains a Q-function using parameters $\theta$ on dataset $\mathcal{D}$ by minimizing temporal difference (TD) error.

$$ $\displaystyle\mathcal{L}_{TD}(\theta)=\mathbb{E}_{(s,a,r,s^{\prime})\sim \mathcal{D}}\left[\left(r+\gamma\max_{a^{\prime}}Q_{\bar{\theta}}(s^{\prime},a ^{\prime})-Q_{\theta}(s,a)\right)^{2}\right]\,,$ (1) $$

where $Q_{\theta}(s,a)$ is the parameterized Q-function, and $\bar{\theta}$ parameterize a target network and is a slow-moving copy of $\theta$.

#### RL for language generation.

Language generation can be viewed as an MDP, where states are sequences of tokens from a finite vocabulary ${\mathcal{V}}$ (Ramamurthy et al., ).
All tokens that the agent initially observes are used as our initial state, $s_{0}=(x_{0},\ldots,x_{m})$, where $x_{i}\in{\mathcal{V}},\forall i\in[m]$.
At timestep $t$, an action $a_{t}\in{\mathcal{V}}$ is some token in the vocabulary.
As long as $a_{t}$ is not a special end-of-sequence <EOS> token, the transition function deterministically appends $a_{t}$ to state $s_{t}$ to form $s_{t+1}$. Otherwise, the agent observes (potentially stochastic) responses from the environment, i.e. utterances by conversational partners in the case of goal-directed dialogue, $o_{t}=(y_{0},\ldots,y_{n})$, which also consist of tokens in the vocabulary; then, the transition function appends both $a_{t}$ and responses $o_{t}$ to state $s_{t}$.
This continues until the last timestep $T$ where we obtain a state $s_{T}$ and the agent receives a deterministic reward $r(s_{T})$.

It becomes clear that a policy $\pi(a|s)$ is a language model that parses all the language tokens seen so far as the state, and computes a distribution over tokens as the next action to take. Recently, RL has been considered for learning policies that are LLMs or VLMs for difficult tasks such as generalist robotic manipulation or dialogue.
Because value learning is very different from traditional next-token prediction, preforming such fine-tuning requires reparameterizing the pretrained language model, such as by adding value heads with independently initialized weights (Snell et al., ).

## 4 Q-Learning via Supervised Fine-Tuning

We will now describe our proposed offline RL algorithm, which we dub Q-learning via Supervised Fine-tuning (Q-SFT). Concretely, instead of training value functions by fitting Q-values to their Bellman backup target via a regression loss, we instead fine-tune directly on the probabilities learned from large-scale pretraining
—like in SFT— via a *weighted cross-entropy* loss, such that the resulting probabilities also capture the desired Q-values.

### 4.1 Learning Values as Probabilities

Recently, large neural networks such as LLMs and VLMs have been successfully trained and fine-tuned on demonstration data using supervised learning.
If we adopt the earlier multi-turn formalism in Section 3 and view these models as agents, such approaches train a policy $\pi_{\phi}(a|s)$ with parameters $\phi$ by minimizing cross-entropy loss:

$$ $\displaystyle\mathcal{L}_{\mathrm{CE}}(\phi)=\mathbb{E}_{(s,a)\sim\mathcal{D}} \left[\log\pi_{\phi}(a\mid s)\right]\,.$ (2) $$

Because the resulting policy approximates the behavior policy $\pi_{\phi}(a|s)\approx\pi_{\beta}(a|s)$, this approach has also been dubbed behavioral cloning (BC). While BC scales well to complex tasks and networks, the resulting policy can only be as good as the behavior policy, which is insufficient when the dataset is not curated from expert demonstrations.

In contrast, Q-learning enables the learned policy to greatly outperform the behavior policy  (Kumar et al., ), by instead having the policy behave according to the estimated Q-values. This can be done via *policy extraction*, such as $\pi(a|s)=\mathbbm{1}[a=\arg\max_{a}^{\prime}Q_{\theta}(s,a^{\prime})]$ or the entropy-regularized variant $\pi(a|s)\propto\exp(Q_{\theta}(s,a))$.
However, as alluded to earlier, the Q-function $Q_{\theta}(s,a)$ cannot be naturally derived from pretrained language models, which output probabilities, and require modifying their architectures as in Figure 1.

Our goal is to provide a way to learn Q-values for multi-turn RL problems with language models such that the Q-function can be initialized from a model pretrained via supervised learning (i.e., maximum likelihood estimation), *without* the need to reinitialize weights or add new heads to represent the Q-values.
An autoregressive sequence model (e.g., a transformer) outputs the probability of each token conditioned on the past history.
In order to avoid adding new heads or reinitializing weights, the Q-values have to also be represented by these same probabilities.
Furthermore, to maximize transfer from pretraining, we would like our proposed *loss function* to also closely resemble the maximum likelihood loss function used for pretraining.

We propose a simple modification to the BC objective in Equation 2. Our modification hinges on the following observation. Let $p_{\theta}(a|s)$ represent the probability of action $a$ under state $s$, and are optimized via the *weighted* cross entropy loss

$$ $\mathcal{L}_{\mathrm{WCE}}(\theta)=\mathbb{E}_{(s,a)\sim\mathcal{D}}\left[w(s, a)\log p_{\theta}(a\mid s)+(1-w(s,a))\log p_{\theta}(a_{d}\mid s)\right]\,,$ $$

where $w(s,a)$ are weights, and $a_{d}$ is some dummy action.
The resulting probabilities that optimize this objective approximate $\widehat{p}_{\theta}(a|s)\approx w(s,a)\pi_{\beta}(a|s)$ for all $a\neq a_{d}$.
Our goal is, via a proper choice of weights, to learn probabilities that are conservative estimates of the true Q-values $\widehat{p}_{\theta}(s,a)\approx Q^{*}(s,a)$.
In order to do so, we require the following assumption on bounded total rewards:

###### Assumption 4.1 .

For any policy $\pi$, we have $\sum_{t=1}^{\infty}\gamma^{t-1}r_{t}\leq 1$.

This assumption has been made by multiple prior works without loss of generality
(Ren et al., ; Kumar et al., ), as rewards can, in theory, be scaled without affecting the optimal policy in the MDP. Furthermore, many tasks of interest, such as dialogue, have sparse rewards, where we observe success or failure only after the conversation has ended.

Following the above observation, let us define the *empirical Bellman probability operator* $\widehat{\mathcal{B}}^{*}$ for transition $(s,a,r,s^{\prime})$ as

$$ $\widehat{\mathcal{B}}^{*}p_{\theta}(a\mid s)=r+\gamma\max_{a^{\prime}}\frac{p_ {\theta}(a^{\prime}\mid s^{\prime})}{\pi_{\beta}(a^{\prime}\mid s^{\prime})}\,.$ $$

Note that this is different from the traditional Bellman operator in that we additionally divide by $\pi_{\beta}$ in the backup. Then, we consider the following weighted cross-entropy loss:

$$ $\displaystyle\mathcal{L}_{\mathrm{WCE}}(\theta)=\mathbb{E}_{(s,a,r,s^{\prime}) \sim\mathcal{D}}\left[\widehat{\mathcal{B}}^{*}p_{\bar{\theta}}(a\mid s)\log p _{\theta}(a\mid s)+\sum_{a^{\prime}\neq a}\frac{1-\widehat{\mathcal{B}}^{*}p_{ \bar{\theta}}(a\mid s)}{|\mathcal{A}|-1}\log p_{\theta}(a^{\prime}\mid s) \right]\,.$ (3) $$

Here, we see that our loss is an instance of weighted cross entropy loss with weights approximately equal to Bellman target values $\widehat{\mathcal{B}}^{*}p_{\bar{\theta}}(a|s)$.
The primary difference is that instead of introducing a dummy action, we equally distribute the leftover weight across the remaining actions. As we will show, this acts as a label-smoothing term that ultimately regularizes the probabilities.
We will show later in this section that in the absence of sampling error, our learned likelihood function $\widehat{p}_{\theta}(a|s)$ satisfies $\pi_{\beta}(a|s)\ Q^{*}(s,a)\leq\widehat{p}_{\theta}(a|s)\leq Q^{*}(s,a)$.
This means that we are able to effectively learn a conservative estimation of the Q-function as a likelihood, without the need for optimizing a potentially unstable and poorly-scaling TD objective.

In addition, because probabilities are modeled directly by existing language models, we do not need to modify the parameterization of such models in order to perform such fine-tuning, i.e. by resetting weights or adding a new head.
Namely, our likelihood function $\widehat{p}_{\theta}(a|s)$ can be directly initialized from the logits of a pretrained LLM or VLM.

###### Assumption 4.1 .

### 4.2 Practical Implementation

Our objective in Equation 3 trains the model such that the predicted token probabilities match their Q-values.
Our final step is to choose how to use these probabilities in the final learned policy.
Prior work performs policy extraction that learns a policy regularized to be similar to the behavior policy (Peng et al., ; Kostrikov et al., ).
Namely, it is well-known that a policy parameterization $\widehat{\pi}(a|s)\propto\pi_{\beta}(a|s)\exp(Q^{*}(s,a))$ is a solution to the constrained optimization problem (Peng et al., ; Brandfonbrener et al., )

$$ $\displaystyle\arg\max_{\pi}\ \mathbb{E}_{s\sim d^{\pi_{\beta}},a\sim\pi}\left[ Q^{*}(s,a)\right]\quad\text{ s.t. }\quad\mathbb{E}_{s\sim d^{\pi_{\beta}}} \left[\mathrm{D}_{\mathrm{KL}}(\pi(\cdot\mid s)\ ||\ \pi_{\beta}(\cdot\mid s)) \right]\leq\varepsilon.$ $$

Recall that we have already approximated $\pi_{\beta}$ by optimizing an $\mathcal{L}_{CE}(\phi)$ over parameters $\phi$.
Therefore, without any further training we can extract a policy whose probabilities of actions $\widehat{\pi}(a\mid s)$ follow
$\widehat{\pi}(a|s)\propto\pi_{\phi}(a|s)\exp\left(p_{\theta}(a|s)\right)$,
an expression of both our learned probabilities $p_{\theta}$ and estimated behavior policy $\pi_{\phi}$.
Note that unlike prior works that explicitly require a policy extraction step with additional training, our policy can be computed at inference-time. Namely, we can express our policy using only the probabilities that we had previously learned.
An overview of our entire method is provided in Algorithm 1.

Figure: Algorithm 1 Q-learning via Supervised Fine-tuning (Q-SFT)

### 4.3 Theoretical Analysis

To simplify exposition, we consider a modified version of our objective in Equation 3, where instead of empirical operator $\widehat{\mathcal{B}}^{*}$, we use the true operator:

$$ $\mathcal{B}^{*}p_{\theta}(a\mid s)=r(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim P( \cdot|s,a)}\left[\max_{a^{\prime}}\frac{p_{\theta}(a^{\prime}\mid s^{\prime})} {\pi_{\beta}(a^{\prime}\mid s^{\prime})}\right]\,.$ $$

Note that it is simple to adapt our analysis to the empirical operator. Namely, it requires obtaining high-probability bounds of the form
$\forall(s,a),\ \left|\widehat{\mathcal{B}}^{*}p_{\theta}(a\mid s)-\mathcal{B}^
{*}p_{\theta}(a\mid s)\right|\leq\frac{C}{\sqrt{n(s,a)}}$
where $C$ is a constant independent of $\gamma$. This kind of inequality commonly arises in analysis of offline RL algorithms  (Kumar et al., ; ).

Our main theoretical result is that our learned probabilities satisfy being conservative estimates of the true value function:

###### Theorem 4.1 .

Let $\widehat{p}_{\theta}$ be the likelihood function that arises from optimizing Equation 3 using the true Bellman likelihood operator. Then, $\widehat{p}_{\theta}$ satisfies

$$ $\displaystyle Q^{*}(s,a)\geq\widehat{p}_{\theta}(s,a)\geq Q^{*}(s,a)\pi_{\beta }(a\mid s)\,,$ $$

for all $s\in\mathcal{D}$ and $a\in\mathcal{A}$ such that $Q^{*}(s,a)\geq\frac{1}{|\mathcal{A}|-1}$.

We defer proof of the theorem to Appendix A.
Note that our probabilities are conservative only over actions that have non-negligible Q-values. In practice, we do not see this as a problem as actions with negligible Q-values will not be chosen anyway.
Overall, we show that while our objective looks very different from traditional value-based RL, our method still learns a conservative value function.

So far, we have shown that theoretically, our algorithm achieves similar theoretical properties as value-based RL methods, even though our objective looks closer to supervised learning.
Next, we will compare our approach to other RL algorithms adapted from supervised learning such as filtered behavior cloning or return-conditioned supervised learning, and show why ours is beneficial.

Filtered behavior cloning. Filtered BC attempts to adapt supervised fine-tuning to non-expert datasets by only training on the top $\rho$-percent of trajectories by reward for $\rho\in[0,1]$. While natural, this harms sample efficiency as our method, like value-based RL methods, is also able to extract meaningful knowledge from low-reward trajectories.

Return-conditioned supervised learning. RCSL remedies the issues with filtered BC by conditioning on reward during training to extract multiple policies rather than just the expert one.
However, we argue that RCSL still fails to learn from low-reward trajectories as well as our approach.
Specifically,
our approach and others that learn value functions can learn from suboptimal trajectories by stitching them into a better policy  (Fu et al., ).
However, Brandfonbrener et al. () showed that RCSL cannot perform stitching in general, which limits its effectiveness at learning from suboptimal data.

###### Theorem 4.1 .

## 5 Experiments

Our method combines aspects of both SFT and value-based RL training, and we therefore compare our method to state-of-the-art methods from both classes, evaluating:

- (1)
Whether our method improves over SFT methods by taking into account and optimizing over a multi-step task reward.
- (2)
Whether our method improves on the stability and performance of previously proposed value-based RL methods for training LLMs and VLMs.
- (3)
Whether our method is better able to benefit from the pretraining of large models than previously proposed multi-turn RL methods.

In this section, we perform a comprehensive empirical evaluation across a suite of different tasks to find positive answers to all the above questions.

### 5.1 Task Descriptions

Figure: Figure 2: Overview of all the evaluated tasks, spanning both text and image inputs. Solving all the tasks effectively requires our algorithm to be able to be used to fine-tune LLMs, VLMs, and even robotics transformer models.
Refer to caption: x3.png

Contrary to many existing applications of RL on language models, such as RLHF (Ouyang et al., ) or DPO (Rafailov et al., ), our proposed algorithm is tailored for offline RL on *multi-step* tasks. Therefore, we consolidate a variety of existing benchmarks where a language model must make sequential decisions, arriving at the following suite of different tasks.

The first set of tasks include language games from the LMRL benchmark (Abdulhai et al., ), which is one of the first benchmarks tailored at evaluating offline RL for language generation.

Chess. This task uses a textual representation of the game of chess. The offline dataset consists of trajectories by Stockfish 15.1 simulating various player strengths as the agent, playing against another Stockfish engine with Elo 1200.
The reward is 1 for a move that results in victory, 0 for a legal move and -1 for an illegal move. Our dataset consists of 625K trajectories of full games, in which the agent achieves an average return of 0.21.

Wordle. In the game of Wordle, the agent is given at most 6 attempts to guess a hidden 5-letter word. After each guess, the agent is told whether each letter in the guessed word is: (1) in the hidden word and in the right position (green), (2) in the hidden word but not in the right position (yellow), or (3) not in the hidden word (gray). The agent receives a reward of -1 after each incorrect guess. The dataset consists of 20K trajectories by a suboptimal heuristic policy that achieves an average return of -4.12, originally collected by Snell et al. ().

Twenty Questions.
The final language task is the dialogue game of twenty questions, where the agent tries to guess what a hidden object is by asking a series of yes-or-no questions. The dataset consists of 100K conversations between an agent that is a guesser and the oracle that chooses the hidden word. The oracle chooses the hidden work uniformly at random from 158 unique objects. The guesser and the oracle are both simulated using GPT3.5 (OpenAI, ), which is prompted to both generate questions and answer them factually.
The agent receives a reward of -1 for each question that is not a correct guess, up to a minimum return of -20. The average return in the dataset is -17.3.

The next evaluation for fine-tuning LLMs as language agents — interactive web-based tasks that require using tools like search.

WebShop. an online shopping website environment where an agent processes unstructured text data (in the form of descriptions crawled from Amazon) to purchase a product given some initial user specifications. At the end, the agent receives a reward between 0 0 and $1$ depending on the similarity between the purchased and ground-truth desired item. The benchmark consists of $12$k initial user instructions, of which we randomly held out 100 for evaluation. With the remaining instructions, we generate a dataset of trajectories where we simulate an suboptimal agent by prompting GPT3.5 with few-shot examples, following the prompts used by  Yao et al. ().

Our method can be applied not only to language models, but also to multimodal models. In the next experiment, we study the performance of our method on vision-based navigation with VLMs.

ALFWorld. This is a popular text-based environment grounded in image observations  (Shridhar et al., ).
In this environment, the agent is tasked with solving one of $6$ different task types, ranging from finding, moving, and manipulating different household objects within an embodied environment of $120$ rooms. At each timestep, the agent observes a textual description of its location and surroundings with an analogous image, and chooses a text action from a set of admissible actions.
In this environment, we sample $10$k trajectories consisting of a random templated task description, and an attempted execution of the task within $30$ timesteps by a prompted GPT3.5 model for data collection.
In the dataset, the agent only successfully accomplishes the task $34\%$ of the time aggregated across all task types.

Finally, we also evaluate our method for training policies outside of language generation.
Robotics is a popular domain in which offline RL has been proven effective for training per-token Q-values for continuous control (Singh et al., ; Chebotar et al., ).
In these experiments, we do not leverage pretrained language models and simply test the effectiveness of the underlying RL algorithm.

Robotic manipulation. We consider the large-scale robotic manipulation control tasks from Singh et al. ().
In the environment, the agent controls a 7-DoF mobile manipulator in front of a countertop surface to perform two types of tasks: pick up an object and place on the trap in front, or grab an object from the drawer. A sparse reward of $1$ is received if the agent accomplishes the task. Following Singh et al. (), we collect $300$k trajectories of randomized, scripted policies performing one of the two task types. The scripted policies achieve roughly $40\%$ success rate.

We show illustrations of all the considered tasks in Figure 2.

### 5.2 Results

**Table 1: Average scores (for language games), and success rates (for ALFWorld tasks) across $100$ independent evaluations. Our method performs best or near-best across the table, and competitively with prompting a much more complex model.**
|  | language games | alfworld |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | Chess | Wordle | 20Q | Pick | Examine | Clean | Heat | Cool | Pick2 |
| ReAct | 0 0 | $-4.96$ | $-13.2$ | $\mathbf{45}$ | $19$ | $17$ | $7$ | $12$ | $\mathbf{24}$ |
| SFT | $0.11$ | $-3.81$ | $-17.3$ | $38$ | $15$ | 0 0 | $11$ | 0 0 | $18$ |
| ILQL | $0.09$ | $-2.08$ | $-14.2$ | $28$ | $7$ | 0 0 | $5$ | $2$ | $15$ |
| Q-SFT (ours) | $\mathbf{0.15}$ | $\mathbf{-2.11}$ | $\mathbf{-13.1}$ | $39$ | $\mathbf{21}$ | $\mathbf{19}$ | $\mathbf{14}$ | $\mathbf{18}$ | $21$ |

Figure: Table 2: Average score across $100$ held-out instructions in WebShop. Our method performs best, even against prompting a much larger model.

The goal of our empirical results is to show positive answers to all the proposed research questions. In order to do so, we evaluate multiple state-of-the-art supervised and value-based RL methods.

#### Evaluation.

We compare our method Q-SFT against three classes of competing algorithms:

*Prompting*: ReAct  (Yao et al., ) is an extension of chain-of-though prompting  (Wei et al., ), where the pretrained language model is prompting to think and reason multiple steps in advance. We use GPT3.5 (OpenAI, ) as the LLM, and GPT4-V (OpenAI, ) as the VLM.

*Supervised learning*: Supervised fine-tuning (SFT) on the offline dataset. In the case of using non-pretrained models, such as in the robotics task, we rename the method as behavior cloning (BC).

*Value-based RL*: Traditional offline RL algorithms that perform Q-learning to learn value functions. We consider several different algorithms, depending on the task at-hand.
In the case of language games and ALFWorld, we evaluate ILQL (Snell et al., ), a popular approach for language generation. In WebShop, we consider an offline variant of ArCHer (Zhou et al., ) that performs best on the task from prior work. Finally, in robotics manipulation, we consider both CQL  (Kumar et al., ) and Q-transformer (QT)  (Chebotar et al., ), which are both popular and achieve state-of-the-art performance in continuous control.

Since the state-of-the-art LLMs and VLMs often only expose inference APIs, we instead train our considered methods on the GPT2-medium LLM, which consists of 345M parameters (Radford et al., ). For ALFWorld, which requires VLMs, we use LLaVA-1.6 model as the pretrained model (Liu et al., ). Finally, for robotics, we use a randomly initialized Transformer architecture modeled after the popular RT-1 model, which processes images and discretizes the action space into tokens  (Brohan et al., ). Note that for the Chess and Wordle tasks, because their state and action space are unlike natural language, we replace the pretrained weights with a random initialization. Therefore, these tasks, like robotics manipulation, only compares the methods in terms of their effectiveness as RL algorithms.

#### Discussion

We report results of our evaluations in Tables 1, 2,and 3.
For fine-tuning LLMs, Table 1 and  2 show that our Q-SFT method outperforms supervised learning, and different value-based RL methods, sometimes outperforming state-of-the-art by almost $30\%$.
Similarly, for VLMs, as shown in Table 1, our approach beats supervised and value-based RL baselines, particularly on hard task types with low success rate in the data.
Our approach is also competitive with state-of-the-art prompting of GPT4-V, which contains about $30\times$ more parameters than the base models using during training.
Finally, in Table 3, we see that even without leveraging any pretraining, our approach is competitive with state-of-the-art, suggesting that our objective is effective for learning values.
This can be attributed to the fact that our underlying objective is more stable to optimize than traditional Q-learning ones, which require regression to non-stationary Bellman target values. Furthermore, in Figure 3, we show the learning curve for our approach compares favorably to QT, showing that our method learns more quickly and achieves better performance in the low-data regime.

Figure: Figure 3: Success rate during initial training on the pick object task of the robotic manipulation benchmark. Though our method achieves similar final performance as Q-transformer, we perform much better on fewer samples.
Refer to caption: x4.png

Finally, we want to answer the last research question. We hypothesize that our approach benefits more from pretraining than existing value-based RL techniques, and verify this with an additional experiment.
Specifically, we consider the 20Q task, and train both ILQL and Q-SFT policies against different models of increasing number of parameters, namely the GPT2-large and GPT2-xl models, which are $2\times$ and $5\times$ larger than GPT2-medium respectively.
We also only train on $10\%$ of the original dataset, so that retaining prior knowledge from pretraining becomes crucially important.
In Figure 4, we show the average return achieved by both methods across the different model sizes. We notice that for larger model sizes, Q-SFT significantly outperforms ILQL, implying a positive answer to the research question that our method retains knowledge acquired during pretraining better than existing value-based RL.

## 6 Discussion

In this paper, we present Q-learning via Supervised Fine-Tuning (Q-SFT), a new offline RL algorithm where Q-values are learned as probabilities in an objective that looks like supervised fine-tuning.
Because of this, our objective can be directly optimized over the logits of pretrained LLMs or VLMs
To our knowledge, this is the first algorithm that can perform value-based RL fine-tuning without requiring any changes to the architecture, such as adding in new value heads.
This has a number of important benefits.
First, our objective is an instance of weighted cross-entropy, which has been shown by prior works to be more stable to train than traditional value-based RL methods that require regression towards non-stationary target values.
More importantly, our algorithm fully leverages the advantage of foundation models such as LLMs or VLMs, as our algorithm starts from the pretrained probabilities, as opposed to randomly-initialized values.
Theoretically, we show that our probabilities are conservative estimates of the true value function. Empirically, we compare our approach against strong supervised and value-based RL baselines on a variety of different tasks requiring LLMs, VLMs, and even robotics transformers.
As future work, we aim to use our approach to also fine-tune vision-language-action (VLA) models, where we expect to see even greater benefit (Kim et al., ).
Another more interesting direction for futher investigation, is whether our method can be adapted to also work online, as many recent works have considered online Rl optimization of language agents (Zhou et al., ).

## Appendix A Theoretical Proofs

Here, we provide a proof of Theorem 4.1. Recall that we are optimizing the following objective:

$$ $\displaystyle\mathcal{L}_{\mathrm{WCE}}(p)=\mathbb{E}_{(s,a,r,s^{\prime})\sim \mathcal{D}}\left[\mathcal{B}^{*}\widehat{p}(a\mid s)\log p(a\mid s)+\sum_{a^{ \prime}\neq a}\frac{1-{\mathcal{B}}^{*}p(a\mid s)}{|\mathcal{A}|-1}\log p(a^{ \prime}\mid s)\right]\,.$ (4) $$

Let us consider iteration $k$ of training. Setting the derivative of Equation 4 to 0 0, we obtain the following expression of $\widehat{p}^{k+1}$ in terms of $\widehat{p}^{k}$:

$$ $\displaystyle\forall s\in\mathcal{D},a\in\mathcal{A},\quad\widehat{p}^{k+1}(a \mid s)=\pi_{\beta}(a\mid s)\mathcal{B}^{*}\widehat{p}^{k}(a\mid s)+\sum_{a^{ \prime}\neq a}\pi_{\beta}(a^{\prime}\mid s)\ \frac{1-\mathcal{B}^{*}\widehat{p }^{k}(a^{\prime}\mid s)}{|\mathcal{A}|-1}\,.$ (5) $$

Lower-bound. We will first show the lower-bound part of Theorem 4.1. Rearranging the above equation, we see that:

$$ $\displaystyle\frac{\widehat{p}^{k+1}(a\mid s)}{\pi_{\beta}(a\mid s)}=\mathcal{ B}^{*}\widehat{p}^{k}(a\mid s)+\sum_{a^{\prime}\neq a}\frac{\pi_{\beta}(a^{ \prime}\mid s)}{\pi_{\beta}(a\mid s)}\ \frac{1-\mathcal{B}^{*}\widehat{p}^{k}( a^{\prime}\mid s)}{|\mathcal{A}|-1}\,.$ (6) $$

Hence, we see that

$$ $\displaystyle\frac{\widehat{p}^{k+1}(a\mid s)}{\pi_{\beta}(a\mid s)}\geq \mathcal{B}^{*}\widehat{p}^{k}(a\mid s)=r(s,a)+\gamma\mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[\max_{a^{\prime}}\frac{\widehat{p}^{k}(a^{\prime}\mid s ^{\prime})}{\pi_{\beta}(a^{\prime}\mid s^{\prime})}\right]\,,$ $$

where we substitute the definition of $\mathcal{B}^{*}$. Finally, taking the fixed point of the above expression yields,

$$ $\displaystyle\frac{\widehat{p}(a\mid s)}{\pi_{\beta}(a\mid s)}\geq Q^{*}(s,a) \Rightarrow\widehat{p}(a\mid s)\geq\pi_{\beta}(a\mid s)Q^{*}(s,a)\,,$ $$

as desired.

Upper-bound. Now, we show the upper-bound part of Theorem 4.1.
Assume that

$$ $\mathcal{B}^{*}\widehat{p}^{k}(a\mid s)\geq\frac{1}{|\mathcal{A}|-1}\,.$ $$

Then, we can solve for the bound:

$$ $\displaystyle\left(1-\pi_{\beta}(a\mid s)\right)\ \mathcal{B}^{*}\widehat{p}^{ k}(a\mid s)-\sum_{a^{\prime}\neq a}\pi_{\beta}(a^{\prime}\mid s)\ \frac{1- \mathcal{B}^{*}\widehat{p}^{k}(a^{\prime}\mid s)}{|\mathcal{A}|-1}$ $\displaystyle=\sum_{a^{\prime}\neq a}\pi_{\beta}(a^{\prime}\mid s)\ \left( \mathcal{B}^{*}\widehat{p}^{k}(a\mid s)-\frac{1-\mathcal{B}^{*}\widehat{p}^{k} (a^{\prime}\mid s)}{|\mathcal{A}|-1}\right)$ $\displaystyle\geq\sum_{a^{\prime}\neq a}\pi_{\beta}(a^{\prime}\mid s)\ \left( \frac{1}{|\mathcal{A}|-1}-\frac{1-\mathcal{B}^{*}\widehat{p}^{k}(a^{\prime} \mid s)}{|\mathcal{A}|-1}\right)\geq 0\,.$ $$

This means that we have,

$$ $\displaystyle\widehat{p}^{k+1}(a\mid s)$ $\displaystyle=\pi_{\beta}(a\mid s)\mathcal{B}^{*}\widehat{p}^{k}(a\mid s)+\sum _{a^{\prime}\neq a}\pi_{\beta}(a^{\prime}\mid s)\ \frac{1-\mathcal{B}^{*} \widehat{p}^{k}(a^{\prime}\mid s)}{|\mathcal{A}|-1}$ $\displaystyle=\mathcal{B}^{*}\widehat{p}^{k}(a\mid s)-\left(1-\pi_{\beta}(a \mid s)\right)\ \mathcal{B}^{*}\widehat{p}^{k}(a\mid s)+\sum_{a^{\prime}\neq a }\pi_{\beta}(a^{\prime}\mid s)\ \frac{1-\mathcal{B}^{*}\widehat{p}^{k}(a^{ \prime}\mid s)}{|\mathcal{A}|-1}$ $\displaystyle\leq\mathcal{B}^{*}\widehat{p}^{k}(a\mid s)\,.$ $$

Hence, we see that

$$ $\displaystyle\frac{\widehat{p}^{k+1}(a\mid s)}{\pi_{\beta}(a\mid s)}\leq\frac{ \mathcal{B}^{*}\widehat{p}^{k}(a\mid s)}{\pi_{\beta}(a\mid s)}=\frac{1}{\pi_{ \beta}(a\mid s)}\left(r(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim P(\cdot|s,a)} \left[\max_{a^{\prime}}\frac{\widehat{p}^{k}(a^{\prime}\mid s^{\prime})}{\pi_{ \beta}(a^{\prime}\mid s^{\prime})}\right]\right).$ $$

Finally, taking the fixed point of the above expression yields the desired $\widehat{p}(a\mid s)\leq Q^{*}(s,a)$. This completes the proof.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - While the theoretical analysis shown in Section 4.3 demonstrates that the Q value learned with Equation (3) has the bounds independent from the hyperparameter \gamma, I wonder what the ablating effect \gamma brings in practice: Will certain configuration of \gamma contribute the learning of a better Q-value estimation? This question has another motivation: In Equation (3), some of the probability mass is allocated to the dummy actions a', and the weight ratio between the actions in the demonstration trajectory and the dummy ones has one degree of freedom determined by \gamma. It would therefore be great to sweep \gamma to study the ablation effect of the weight ratio.
   - And by the way, how is the \gamma set in the demonstrated experiments?

- According to Equation (3), learning the Q-value function requires a learned behavior policy \pi_\beta in prior. How will the quality of the behavior policy impact the learning of the Q-values, and the performance of the ultimate policy? While I appreciate the experiments in Figure 4 where only 10% of the training data are used to get the \pi_\beta, I wonder how the performance grows with the quality of the initial behavior policy. For example, if the initial behavior policy is near perfect, will it be that the ultimate policy, although weighted with the learned Q-values, performs similarly with the initial one?
soundness: 4
strengths: - The work targets at the multi-turn RL scenario for LLMs, which is vital for the development of complex reasoning and agentic use in LLMs.
- The proposed algorithm is backuped with a solid theoretical motivation.
- The experimenting scenarios have a broad coverage, and the performance is well-demonstrated.
summary: This paper proposed Q-SFT, a method that integrates Q-learning with SFT for LLMs in the multi-turn RL setting. To exploit both the representation and the logits in LLMs learned in the pretraining stage and bypass the poor scaling effect of TD-style objectives, the authors proposed a novel weighted cross-entropy loss that embeds the learning of the Q-value into the weights. In this way, the logits prior in the pretrained LLMs can be leveraged, and no new head is needed to learn the Q-values. Theoretical analysis has shown the guarantee of the new learning objective leading to a conservative estimation of the value function. Experiments on several scenarios have demonstrated the effectiveness and efficiency of the proposed method.
weaknesses: Please see the Questions below.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 1
questions: Please see the previous section.

Response to the rebuttal: (posted here since the rebuttal ended)

Thanks for all the clarification and discussion -- I raised my scores accordingly. It would be very valuable to revise the paper to address the points we discussed, especially the clear definition on step vs turn, the potential to scale better than existing RL approaches, the setup and comparable work of SFT baselines, motivation to alleviate the downsides of using RL, etc.. The rebuttal was helpful to better understand this work, and it would definitely be better if readers only need to read the paper itself to sufficiently appreciate this work.
soundness: 3
strengths: - Using probability prediction in the Q-learning objective is cool, and it sounds particularly beneficial for large vocabulary/action space. 
- The implementation of the method is straightforward: learning two models using different objectives, and use both at inference time
summary: This paper proposes a method that leverages the LLM logits learned via supervised fine-tuning to approximate the Q-learning objective. Authors argue that traditional Q-learning setup suffers from discarding the logit head, and learning a new head for Q values, so the proposed method directly translates the learned logits to Q values, with theoretical analysis on its bound and approximation. Experiments include text-based games, image-based navigation and robotic manipulation, and show some benefits of the proposed method over the prior value-based RL method.
weaknesses: - I found the pitch very confusing: it emphasizes multi-turn problems like dialogue, but the actual math formulation treats each turn as a single token in an utterance (i.e., the action space is the vocabulary space). The generation process is to generate an utterance token by token, which is conventional, and it’s confusing where the multi-turn setup plays a role other than using this generative model repeatedly to generate different utterances at different parts of the dialogue. In other words, the value function here is not the reward per utterance as stated in the intro (line 074), and question answering is also not a single-step problem as the answer is also generated token by token.
- Experimental setup seems simplified in terms of 1) for language-based tasks, the game nature makes the effective vocabulary very small and unambiguous, 2) the well-defined reward function setup in text games do not translate into practical problems, 3) why no supervised learning results from prior work as a direct comparison?; 4) most of the datasets already have RL method outperform supervised learning which is contradictory to the intro “offline RL fine-tuning of LLMs does not reliably outperform supervised fine-tuning”. Therefore, it’s unclear whether the improvement brought by the proposed method should be contrasted to what exactly.
- The writing can be further improved in terms of clarity, accuracy, and fixing typos. A few examples below: 
   - Table 1: number from the proposed method is bold in the Wordle column, but it’s not the best performance method
   - Define h used in line 162
   - Prime should be on a, instead of argmax in line 218
   - Consistently use “RL” instead of having “RI” occurred 
   - Citation should be within a pair of parentheses: line 037, 039

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: - Some typos: L39 citation, L84 sentence, L194 “performing”, Fig2 top left, 
- Eq2 should have a negative sign? Also tab1 -2.08 should be bold, not -2.11. 
- The figure 1 was bit too small and hard to read.
- L153 defines RL as MDP, but what about POMDP?
- About the theoretical guarantee, how good is the lower bound? A probability from a policy can be quite low for large vocabulary, in which case then the bound is not that tight?
- How many models needed to be loaded during training? There is the main model, also a target model and the original model. So is that 3? How does that affect memory usage?
- Given that some baselines use different base models, it would be nice to put them in the table. This will help if they perform better because their base model is stronger or not.
- In fig3, the error bars correspond to different seeds?
soundness: 3
strengths: - Novelty: the method seems to be quite novel. The way it works is quite different from a typical Q-learning where L2 loss is used. Instead it relies on a cross-entropy loss that is modified. Given the importance of RLHF in LLMs and the need of multi-step interaction, a study like this can be impactful and important.
- The experiments are quite diverse, ranging from LLM-based text tasks to VLM-based robotic manipulation. It is also compared to a diverse set of baselines, which helps to solidify the claims.
- There is a proof about a theoretical guarantee, which is always nice to have. The method also leads to an improvement in performance in practice too, where the gap increasing as the model scales. Given the lack of value-based approaches in LLMs, it would be interesting to see if this method will adopted widely.
- Paper is well written and easy to follow.
summary: This paper proposes a new RL training method for LLMs that take advantage of the pretraining. It is based on Q-learning, but instead of attaching a new value head, the method uses the existing token probabilities as value function. There is some theoretical guarantees and diverse experimental results involving multi-steps interactions.
weaknesses: - I felt like the writing of the experimental results were a bit too short. It would have been good to have more ablation studies to understand why the method works, and in which situations it is more optimal. For example, since the main motivation is that the value head is not initialized from scratch, what will happen if we do that but keep the loss the same. Also there is not much on the training details, such as the number of training steps etc. I think the other parts can be condensed to make space. 
- Another important part that was too short was the method itself. The method is introduced between L230-261, which is about only half page. It doesn’t give enough explanation about why this method should work, so it was quite hard to understand it. For example, can you give more insight into why the proposed method should be more stable? 
- Reuse of the pretrained weights is emphasized as the main motivation, but some of the tasks actually train the model from scratch and others are not in natural language, which is a bit conflicting. There is still an improvement when trained from scratch, which is good, but there is a lack of explanation why it works better.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: .
soundness: 3
strengths: - The motivation is clear: the instability of the Q-learning objective is a well-known challenge in the RL literature.
- The proposed method is straightforward and demonstrates strong empirical performance, with improvements in sample efficiency and scalability over previous methods.
- The evaluation covers a wide range of benchmarks.
summary: This paper proposes Q-SFT, a new approach that enables fine-tuning of LLMs and VLMs within an offline RL framework but using a supervised learning-style loss function. The method matches or outperforms previous approaches across a range of LM fine-tuning benchmarks.
weaknesses: It seems there are some limitations to the proposed approach that are not addressed in the paper.
- During inference, the approach requires double the number of forward passes—one for the original pretrained LM and one for the fine-tuned LM. This can be challenging in practical scenarios where real-time decision-making is essential, such as robotic manipulation.
- In Theorem 4.1., the lower bound for the learned p_theta is Q* multiplied by pi_beta. If pi_beta is small, the resulting pi_theta can be overly conservative compared to the true Q*.

Meta-review: This paper proposes a new way to do RL in multi-turn dialogue setting where the use of RL is less common compared to the vast success of RL in single-turn case. The proposed approach frames Q-learning as a cross-entropy problem with suitably selected weights such that the trained policy can be used to recover conservative bounds on the Q-value. 

While I think the mathematical reasoning here is not important, the reviewers appreciated the mathematical analysis. However, upon going over the mathematical derivations in the paper I noticed that certain steps were not justified. While equation 5 is correct, the next equation isn't justified. It along with the WCE loss assumes that $\mathcal{B}^\star p_\theta$ is in $[0, 1]$ but it isn't clear why this is the case even with assumption 4.1. In particular, if for some $a'$ we have $\pi_\beta(a' \mid s') = 0$, then for a uniform $p_\theta$ we have $\mathcal{B}^\star p_\theta = \infty$. Note that the public comment also points this out. Now, if we assume $\mathcal{B}^\star p_\theta \in [0, 1]$, then we get can complete the proof. These steps are missing in the proof but are as follows:

$\alpha^{(k+1)}(s, a) \ge r(s, a) + \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} \alpha^{(k)}(s', a')]$,

where $\alpha^{(t)}(s, a) = \frac{p^{(t)}(a \mid s)}{\pi_\beta(a \mid s)}$ for any $t \in \mathbb{N}$. Now if you assume convergence that $\lim_{t\rightarrow \infty} \alpha^{(t)}(s, a) = \alpha(s, a)$, we get:

$\alpha(s, a) \ge r(s, a) + \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} \alpha(s', a')]$,

using Bellman optimality, 
$Q^\star(s, a) = r(s, a) + \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} Q^\star(s', a')]$

and defining $\delta(s, a) = \alpha(s, a) - Q^\star(s, a)$ and $\delta = \inf_{s, a}\delta(s, a)$, one gets:

$\delta(s, a) \ge \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} \alpha(s', a') - \max_{a'} Q^\star(s', a')] \ge \gamma E_{s' \sim P(. | s, a)}[ \min_{a'} \left( \alpha(s', a') - Q^\star(s', a')\right)] \ge \gamma \delta$,

where the last line uses $\max_x f(x) - \max_x g(x) \ge \min_x (f(x) - g(x))$. Finally, taking inf on LHS means that $\delta \ge \gamma \delta$, which implies $\delta \ge 0$ if $\gamma \in (0, 1)$. Or, $\alpha(s, a) \ge Q^\star(s, a)$. I'd request authors to add justification for why $B^\star p \in (0, 1)$ and these missing steps.

Strengths:
1. Method is simple
2. Experiments are presented for several domains
3. Multi-turn LLM reasoning is an important topic and the use of RL for this has been somewhat less successful. This approach proposes a solution for this.
4. Reviewers appreciated the theoretical result but see below.

Weakness:
1. The theoretical results are at best unjustified and at worst potentially incorrect
2. Requires 3 copies of the model. This makes it expensive in practice as even 2x increase in memory is costly and makes a given approach undesirable. This is important especially given the motivation of the authors to make the Q-learning approach user-friendly.
3. The reviewer ChGQ requested results where the value head is randomly initialized but the authors did not include this study in response. I understand that experiments can be hard to complete during the rebuttal period but this one would have been nice to see.

Overall, my biggest concern is the mathematical foundation of the approach. However, given that the approach works well empirically and is simple and should easy to try, I am leaning towards a weak accept.

Decision: Accept (Poster)