[figure]style=plain,subcapbesideposition=center [table]capposition=above

# SelfCheck: using LLMs to zero-shot check their own step-by-step reasoning

Ning Miao1* Yee Whye Teh1 Tom Rainforth1 1Department of Statistics, University
of Oxford. *Email<EMAIL_ADDRESS>
###### Abstract

The recent progress in large language models (LLMs), especially the invention
of chain-of-thoughts (CoT) prompting, makes it possible to solve reasoning
problems. However, even the strongest LLMs are still struggling with more
complicated problems that require non-linear thinking and multi-step
reasoning. In this work, we explore whether LLMs have the ability to recognize
their own errors, without resorting to external resources. In particular, we
investigate whether they can be used to identify individual errors within a
step-by-step reasoning. To this end, we propose a zero-shot verification
scheme to recognize such errors. We then use this verification scheme to
improve question-answering performance, by using it to perform weighted voting
on different generated answers. We test the method on three math
datasets—GSM8K, MathQA, and MATH—and find that it successfully recognizes
errors and, in turn, increases final predictive performance.

##  1 Introduction

Recent years have witnessed dramatic changes in the areas of NLP and AI
brought by large language models. From GPT-3 (Brown et al., 2020), PaLM
(Chowdhery et al., 2022), Llama (Touvron et al., 2023) to GPT-4 (OpenAI, 2023)
and PaLM-2 (Google, 2023), the increasing model sizes and exploding amount of
training data have empowered LLMs to achieve human-level performance on a
large range of tasks, including summarization, translation, question answering
and basic mathematical reasoning.

However, the performance of even the largest LLMs on complex reasoning
problems is still unsatisfactory. For example, GPT-4 only correctly answers
42.5% of problems in the MATH dataset, which is far below human level. Such
problems require careful multi-step reasoning to solve, where LLMs are prone
to make mistakes. Namely, to solve such complex reasoning problems, LLMs need
to generate many reasoning steps. Even though their error rate on a single
step is low, the probability of LLMs generating at least one erroneous step
can be fairly high, which makes final predictions less accurate.

It is thus necessary to check for errors in LLMs’ generated reasoning in order
to filter out incorrect predictions and provide an estimation of confidence
for a prediction. Recently, Cobbe et al. (2021) finetune GPT-3 to recognize
errors in the whole reasoning process, and Li et al. (2022) train a BERT model
(He et al., 2020) to check the correctness of each reasoning step. Ling et al.
(2023) eliminate the need to finetune by performing few-shot checking on a
special type of reasoning chains called Natural Program. To check a reasoning
chain, most previous methods look for error patterns that match those in
training data or instructions. However, there are an infinite number of error
patterns, some of which are quite covert and deceptive. As a result, these
methods do not work well without massive training data or examples.

To address this shortfall, we introduce SelfCheck, a zero-shot checking scheme
for LLMs, which is based on step regeneration. SelfCheck works as a step-by-
step checker, such that it individually checks each step in the LLMs reasoning
process based on the available context. For each reasoning step in an LLM’s
output reasoning, SelfCheck first collects related information to form a
simpler context and then checks it by step-regeneration. Then it integrates
the checking results of the individual steps to form a confidence score for
the whole solution. Using confidence scores as weights to vote among multiple
solutions for the same question, SelfCheck in turn allows us to improve
question answering accuracy by focusing on the most accurate answer. SelfCheck
mimics the procedure of human beings double-checking their solutions and
completely bypasses the problem of pattern-matching style checking. Its design
is based on the key observation that LLMs are trained on far more generation
tasks than checking tasks. SelfCheck works in a fully zero-shot manner to keep
generalizability to different LLMs and tasks as well as to reduce the
influence of prompt engineering on experimental results. Users of SelfCheck
can easily change any of its components to a few-shot, or finetuned version to
achieve better results on specific tasks.

We evaluate SelfCheck on three math tasks, namely GSM8K, MathQA, and MATH, in
increasing order of difficulty. For all datasets, we find that using SelfCheck
achieves a significant increase in prediction accuracies compared with
majority voting. We also see that SelfCheck provides an accurate confidence
estimation for LLM’s solutions, which decreases the proportion of incorrect
solutions by 9%, 22.8%, and 16.2% on three datasets, after filtering out low-
confidence solutions. We also perform an extensive analysis on the effect of
per-question sample sizes and some critical design choices of step-checking
methods, which, we hope, could provide some inspiration for future works on
checking methods. Our code is available at
https://github.com/NingMiao/SelfCheck.

To summarize, the main contributions of this paper are as follows:

  * •

We propose SelfCheck, an effective new schema for checking step-by-step
reasoning in LLMs, which provides a better confidence estimation and in turn
improves predictive accuracy.

  * •

We demonstrate that LLMs can do effective self-verification without resorting
to any external resources.

##  2 Method: SelfCheck

In this section, we propose SelfCheck, which is a scheme to check step-by-step
solutions from an LLM by itself. For simplicity of illustration, we restrict
our discussion to math problems, but SelfCheck can be easily applied to other
domains that require multi-step reasoning.

We call the LLM that generates the original step-by-step solution the
generator, which can be a zero or few-shot LLM. The input to the generator is
a question and its output is a multi-step solution in the form of [Step 1,
Step 2, …, Step N]. SelfCheck first checks the correctness of each step of the
solution and then integrate stepwise checking results to form a confidence
score for the whole solution. In the following, we introduce how we perform
step checking and the design of the integration function. Please see Figure 1
for a real example of how SelfCheck works.

Figure 1: A real example illustrating how SelfCheck works. The blocks show the
question, the outputs of the generator, and the checker, respectively. To
check the correctness of a certain step (step 4 in the example) in a multi-
step reasoning procedure, SelfCheck goes through 4 stages before making a
decision, which is elaborated in Section 2.1. Then the integration function
combines the checking results of all steps to form a confidence score for the
whole solution.

###  2.1 Step checking

To check individual steps of the reasoning process, the first thing we should
notice is that the correctness of each step is highly dependent on its
context, namely the question and previous steps in the solution. As a result,
it is not possible to check a single step without context. For example, we
usually need to refer to previous steps for the definition of variables and
the meaning of specific numbers. The target of the step checking in SelfCheck
is thus to check the conditional correctness of each step based on the
provided context. That is, we only care about catching errors at the current
step, and can assume all information from its context to be correct.

A simple idea to achieve this would be feeding the current step as well as all
its context to an LLM and directly asking the LLM to ‘check the correctness of
the step’. However, in practice, we find that this task is too difficult for
the LLM to do effectively. This is first because the checker needs to
understand the key content in the step and then collect all related
information from the context, before actually checking for its correctness.
Moreover, ‘checking’ is not a common task in the training corpus of most LLMs,
which means LLMs do not perform well even if we exemplify in detail how to do
step-checking in the prompt.

To solve this problem, SelfCheck decomposes the checking task into 4 stages:
target extraction, information collection, step regeneration, and result
comparison. SelfCheck uses the LLM to execute each stage successively and uses
the conclusion from the result comparison to predict the correctness of the
original step. The main idea behind the decomposition is to make the LLM focus
on an easier task that it is used to at each stage, which can be seen as a way
of introducing inductive biases into LLMs. Specifically, we prompt the LLM to
first figure out the target of the current step and what information it uses
to achieve the target. Then we ask the LLM to re-achieve the target using only
the collected information. The clear description of the target and the
simplified context make the regeneration stage less challenging for an LLM. As
a result, its output should be more reliable to serve as a reference. The last
step would be comparing the original step with the regeneration output. If
their main conclusions match/mismatch, we think the original step to be
correct/wrong. In the following, we describe each of the subtasks and our
specific instructions to the LLM.

#### Target extraction

To check a step (for example, step 4 in Figure 1), we first need to figure out
what the step is trying to do. Without a specific target, the regeneration
stage would proceed in a random direction, making it impossible to serve as a
reference to the original step. We thus use the LLM itself to extract the
target of a step $i$ using the question and all previous steps (step 0-3 in
Figure 1) with the following prompt:

> The following is a part of the solution to the problem [question]: [Step 0,
> Step 1,…, Step i].
> What specific action does the step [Step i] take? Please give a brief answer
> using a single sentence and do not copy the steps.

During execution, we copy the question and steps into [question] and [step] to
form the actual input to the LLM. The reason for requesting a brief and
single-sentence answer is to avoid getting too much detailed information,
especially specific equations and values, which can be wrong and,
consequently, mislead the regeneration stage.

#### Information collection

To relieve the pressure of the regeneration stage and avoid unrelated
information from affecting the result, we filter out information that is not
directly related to the current step. Specifically, we ask the LLM to select
useful items from the question and all previous items with the following
prompt:

> This is a math question: [question].
> The following is information extracted from the question:
> Information 0: [Information 0]
> Information 1: [Information 1]
> …
> The following are the first a few steps in a solution to the problem:
> Step 0: [Step 0]
> Step 1: [Step 1]
> …
> Step i-1: [Step i-1]
> Which previous steps or information does the next step [Step i] directly
> follow from?

[Information j] is simply the j-th sentence in the question. After retrieving
the output, we extract step or information ids by regular expression. For
example in Figure 1, the current step requires step 2 and 3 and no information
from the question as context. The selected steps and information are then fed
into the regeneration stage.

#### Step regeneration

Given the target and necessary information of the step, we can now ask the LLM
to regenerate it. Because the step is usually a small jump from previous
conclusions, and the information collection stage has already filtered out
irrelevant information, we can usually trust regeneration results. The prompt
of this stage is

> We are in the process of solving a math problem.
> We have some information from the problem:
> Information 0: [Information $I_{0}$]
> Information 1: [Information $I_{1}$]
> …
> The following are some previous steps:
> Step 0: [Step $S_{0}$]
> Step 1: [Step $S_{1}$]
> …
> The target for the next step is: [Target].
> Please try to achieve the target with the information from the problem or
> previous steps.

[Target] is the output from the target extraction stage and [Information] and
[Step] are selected by the information collection stage.

#### Result comparison

The last step is to compare results from the regeneration stage and the
original step with the following prompt.

> The following are 2 solutions to a math problem.
> Solution 1: [Regeneration output]
> Solution 2: [Step i]
> Compare the key points from both solutions step by step and then check
> whether Solution 1 "supports", "contradicts" or "is not directly related to"
> the conclusion in Solution 2. Pay special attention to the difference in
> numbers.

If the regeneration output "supports" or "contradicts" the original step, we
can conclude that the original step should be correct or wrong. Sometimes, the
correctness of the original step cannot be directly inferred from the
regeneration output. For example, when the target is to simplify an equation,
the original step divides all terms by their common devisor while the
regeneration step might combine like terms. In such cases, we are not sure
about the correctness of the original step, which makes ‘is not directly
related to’ the third output choice for SelfCheck.

###  2.2 Results integration

After getting a checking result for each step, we need an integration function
$\phi$ to give a prediction of the overall correctness of the solution. The
input of $\phi$ should be a vector in the form of $[w_{0},w_{1},...,w_{n}]$,
where each item $w_{i}$ represents the step checking result for step $i$. For
simplicity, we assume $w_{i}\in\\{-1,0,1\\}$, where $-1,0,1$ represents
‘contradict’, ‘is not directly related to’ and ‘support’ in the output of the
result comparison stage. Specifically,

$\displaystyle\phi([w_{0},w_{1},...,w_{n}])=2*\text{Sigmoid}\left(-\lambda_{-1}\sum_{i=0}^{n}\mathbbm{1}_{w_{i}=-1}-\lambda_{0}\sum_{i=0}^{n}\mathbbm{1}_{w_{i}=0}\right),$
(1)

where $\lambda_{-1}$ and $\lambda_{0}$ are two non-negative hyperparameters,
chosen on a validation set. We call the output of $\phi$ the confidence score,
which is a float number between $0$ and $1$. It can be directly used as the
weight for voting or binarized by setting a threshold $t\in[0,1]$ to form a
binary checking result for a solution. We take negative steps (-1s) and
neutral steps (0s) into consideration in Equation 1 because the occurrance of
a (possibly) incorrect step could make a reasoning chain to be incorrect. We
ignore 0s, which is less critical because a correct step usually keeps but not
changes the correctness of a reasoning chain.

##  3 Related Work

How to automatically check the correctness of a sequence of reasoning steps is
a long-standing question. The following are 3 main research directions on it.

#### Training/finetuning a verifier

The most straightforward idea would be training a separate verifier model to
check the solution from the generator. For example, Cobbe et al. (2021)
finetune a GPT-3 on GSM8K to predict the correctness of a solution as a whole.
Li et al. (2022) take steps into consideration and train a binary
deberta-v3-large (He et al., 2020) classifier to predict the probability of
correctness for the solution and each step. Paul et al. (2023) train a critic
model to provide structured feedback on intermediate reasoning steps. More
recently, Lightman et al. (2023) build a large dataset, called PRM800K, which
contains labels for step-wise correctness by crowdsourcing. They finetune a
GPT-4 on PRM800K and achieve good results compared with majority voting.
However, finetuning is not always possible due to the limitation of budget and
time. For example, Cobbe et al. (2021) and Lightman et al. (2023) require
finetuning GPT-3 (175B) and GPT-4 for the best performance. Moreover,
finetuning or training on specific datasets or focusing on generating a
specific type of feedback heavily limits the generalisability of the verifier,
making them impossible to be applied to other domains. In comparison,
SelfCheck is a fully zero-shot method, which requires no finetuning and zero
training data. As a result, SelfCheck can be immediately deployed to work on
multiple domains.

#### Verification with external resources

Another chain of research focuses on using external resources to boost
faithfulness or verify LLM results. Lyu et al. (2023) first translate a
question into a symbolic reasoning chain by an LLM and solve the problem by a
deterministic solver. Peng et al. (2023) introduce an external database to
check for incorrect knowledge in LLM outputs. However, these methods are
heavily limited by the ability of external tools. For example, deterministic
symbolic solvers usually have strict constraints on the problem that they can
solve. And Peng et al. (2023) can only check for factoid mistakes.

#### Few-shot verification

Other works show the potential of using few-shot LLMs to correct errors.
Bubeck et al. (2023) find that GPT-4 can serve as a judge to compare LLM
outputs with reference answers. Weng et al. (2022) verify the whole solution
by backward prediction. Specifically, they mask some information in the
question, and make the solution to predict the masked information. If the
masked information is recovered correctly, they label the solution as correct,
and otherwise incorrect. However, this method is based on the hypothesis that
a correct solution should keep every bit of information in the question while
an incorrect one should not. It only holds for some very simple arithmetic or
logic problems, making it unsuitable for more complex tasks. The verifier in
Ling et al. (2023) relies on a specific solution format called Natural Program
to check errors for each step. However, besides only being able to check
solutions in a specific semi-symbolic form, they report worse performance
compared with majority voting. Most of these methods directly ask LLMs to
check reasoning chains, which usually leads to unsatisfactory results because
checking is not a common task that LLMs are trained on. SelfCheck bypasses
this problem by replacing checking with regeneration, whose benefits are shown
clearly in Section 5.2.

##  4 Experiments

We now run experiments on three math-reasoning datasets to evaluate
SelfCheck’s effectiveness in checking multi-step reasoning and improving
generation accuracies.

#### Datasets

GSM8K (Cobbe et al., 2021), MathQA (Amini et al., 2019), and MATH (Hendrycks
et al., 2021) are three large-scale datasets consisting of math problems on
primary school, middle school, and competition levels, containing 1319, 2985,
and 5000 test samples, respectively. For GSM8K and MathQA, we evaluate our
method on their whole test sets while tuning hyper-parameters on a small
subset of their training set. Because of the limitation of resources, we use a
subset of
MATH111https://github.com/lz1oceani/verify_cot/tree/main/results/chatgpt3.5/natural_program/MATH_np.json
test set from Ling et al. (2023). Besides the level of difficulty, the three
datasets differ from each other in the following aspects. Firstly, MathQA
provides 5 options for each problem, while GSM8K and MATH have no options.
Secondly, GSM8K only has arithmetic problems, while MathQA and MATH contain
more diverse problems in geometry, physics, probability, or algebra.

#### LLMs

There are many choices of LLMs, including large-scale ones such as GPTs (Brown
et al., 2020), PaLMs (Chowdhery et al., 2022) and Claude (Anthropic, ), as
well as open-source ones such as LLaMA (Touvron et al., 2023) and FLAN (Wei et
al., 2021). However, because of reasons such as access, budget, and
computation resources, it is only possible for us to run experiments with the
GPT-3.5 API. We could also set up small open-source LLMs such as LLaMA-13B or
FLAN-11B on our own server. However, according to Fu et al. (2023), their
accuracies are below 20% even on the simplest GSM8K dataset, which makes it
meaningless to evaluate on them.

#### Baselines

We use majority voting (or self-consistency decoding (Wang et al., 2022) in
the context of CoT reasoning) as our main baseline following Ling et al.
(2023) and Lightman et al. (2023). Despite its simplicity, majority voting is
quite a strong baseline. Most zero or few-shot methods report similar or even
worse results compared with it (Weng et al., 2022, Ling et al., 2023). It is
impossible to compare with finetuning based methods such as Lightman et al.
(2023), where we have no access to their finetuned models.

(a)  (b)  (c)  (a) GSM8K

(b) MathQA (c) $\text{MATH}^{*}$

Figure 2: SelfCheck achieves higher accuracies compared with the majority-
voting baseline for all sample numbers.

Figure 2 shows the performance gain using confidence scores from SelfCheck to
do weighted voting compared with vanilla majority voting. The upper row shows
that accuracies of SelfCheck and majority voting have the same increasing
tendency as the number of samples increases, which is a result of variance
reduction. We can also see that the yellow lines are always above the blue
ones, which is a result of SelfCheck giving higher weights to more correct
solutions. The $\Delta$-Accuracy plots in the lower row show that the
improvement is significant for all sample numbers, but the gap is higher for
smaller sample numbers. We leave analysis of this phenomenon and results for
higher sample numbers on a smaller dataset in Section 5.1.

(a)  (b)  (c)  (a) GSM8K

(b) MathQA (c) $\text{MATH}^{*}$

Figure 3: Checking accuracies on reasoning process for different threshold
$t$. ‘pred $+$ in real $+$’ is the ratio of samples predicted as correct in
all real correct samples. Similarly, ‘pred $-$ in real $-$’ is the ratio of
samples labeled as wrong in all real incorrect samples. The lower figures show
the balanced average checking accuracy on three datasets.

(a)  (b)  (c)

Figure 4: When increasing the threshold $t$, the confidence of the predicted
correct samples to be real correct increases for GSM8K (67.5%->76.5%), MathQA
(59.4%->82.2%) and MATH (34.6%->50.8%).

Besides serving as a weight calculator to improve the performance of voting,
SelfCheck can also predict the correctness of a single solution. To do so, the
simplest idea would be to set a threshold $t$ to the confidence score. Any
solution with a score above the threshold would be classified as correct or
otherwise wrong. The results for different $t$ are shown in Figure 3. For a
more detailed analysis, we split the solutions into 2 sets, which contain real
correct and real wrong solutions, respectively. The upper row shows the
checking accuracies on each set. For small $t$, every solution will be
classified as correct, so the accuracies for real correct samples
$\text{Acc}_{\text{c}}$ would be $1.0$ while accuracies for real wrong samples
$\text{Acc}_{\text{w}}$ are $0.0$. When $t$ increases, the
$\text{Acc}_{\text{w}}$ increases and $\text{Acc}_{\text{w}}$ drops. The lower
row shows $\text{Acc}_{\text{m}}$ the average of $\text{Acc}_{\text{w}}$ and
$\text{Acc}_{\text{w}}$, which is a indicator of the overall performance. For
different datasets, $\text{Acc}_{\text{m}}$ ranges from 60% to 65%. It is of
course not optimal. However, it is enough to show the effectiveness of
SelfCheck schema. Users can easily add more examples, finetune the LLM for
different stages or combine SelfCheck with external tools, such as code
intepreter to further improve performance.

From another perspective, Figure 4 shows that for large $t$, the proportion of
correct solutions (real $+$) in all verified solutions (pred $+$) is
significantly higher than the case when $t=0$. When $t=0$, all solutions are
labeled as correct, which equals the case without checking. Though this is at
the cost of erroneously labeling about half of real correct samples as
incorrect, it is an important feature for high-stake scenarios, where the risk
of generating an incorrect solution can be much higher than rerunning the
generation process.

##  5 Analysis

In this section, we explore several interesting questions related to SelfCheck
both empirically and theoretically. Limited by budget and time, all
experiments in this section are on a small subset of MathQA test set with 100
randomly selected questions.

###  5.1 Larger sample size per question?

In Figure 2, we can easily see that the performance gain of SelfCheck
decreases as sample numbers increase from 2 to 10. However, it remains a
question what will happen for larger sample numbers. If the gap converges to
0, it means SelfCheck is only a variance reduction method for majority voting,
whose function can be replaced by simply increasing sample numbers for
majority voting. Interestingly, we find in Figure 6 that though the accuracy
of majority voting starts vibrating for sample numbers between 10 and 50, the
accuracy of SelfCheck increases by about 5%. As shown in the lower figure in
Figure 6, the performance gap exhibits a falling-rising pattern. To understand
this pattern, we first need to decompose the prediction error of LLMs into
variance and bias.

(a)  (b)

Figure 5: Generation accuracies after we increase #sample to 50.
$\text{p-value}<0.01$ for $\text{\\#samples}>30$.

#### Variance

For some questions, the probability of an LLM giving the correct answer is
higher than any other answers, but a single run of LLM can still make
mistakes, which is a result of the variance. By the law of large numbers, if
we run the LLM infinite times, we are guaranteed to get the correct answer
simply by majority voting. If we assume the probability of the correct answer
to be $p$, and the probability of the most probable wrong answer to be $q$
($<p$), some simple calculation shows that the probability of getting an
incorrect answer $P_{w}$ after voting is exponentially decreasing with sample
size $n$, which is upper bounded by
$\lceil\frac{1-p}{q}\rceil(\frac{q}{p})^{\lceil\frac{n}{2}\rceil}$. As a
result, when $n>10$, the probability of making mistakes in this case is
already very low, so continuing to increase $n$ would not significantly
increase the accuracy of majority voting.

#### Bias

For some other questions, the probability of the correct answer is not the
highest among all possible answers, which reflects the bias of the LLM. In
such cases, very large $n$ makes majority voting perform worse by always
outputting the incorrect answer. SelfCheck corrects the bias of LLMs by giving
higher weights to real correct solutions and lower weights to incorrect
answers. As $n$ increases, the probability of LLM generating at least 1 real
correct answer keeps increasing, which can be recognized and selected by
SelfCheck. This explains the large accuracy increase of SelfCheck for large
$n$ and shows that LLMs can correct their own biases, without any external
supervision.

###  5.2 Further comparisons

In order to figure out the effect of several critical design choices for
SelfCheck, we compare SelfCheck with some of its variants. Figure 6 and Table
1 show generation accuracies and balanced checking accuracies of each variant.
Table 1 shows the difficulty to achieve high checking accuracies for all
methods. Nonetheless, any methods we considered with a checking accuracy $\geq
57.2\%$ are helpful to prediction performance.

#### Global vs step checking

Figure 6: Generation accuracies for different variants of SelfCheck.

The first question is can we simply ask an LLM to perform a global checking on
the whole solution without checking step by step. To answer it, we prompt the
LLM to perform global checking with the following instruction:

> The following is a question and a solution to it from a student. Carefully
> check whether the solution is correct step by step. End your response with
> your conclusion that starts with "Correct", "Wrong" or "Not Sure".
> Question: [Question]
> Solution: [Step 0, Step 1,…, Step n]

However, we find the global checker outputs "correct" most of the time and
rarely recognizes an error. Consequently, its performance in Figure 6 is very
similar to majority voting and in Table 1 its accuracy is only higher than
random guess (50%). Consequently, LLMs cannot deal with the difficulty of
global checking, which makes step checking necessary.

#### Multiple stages vs Single stage

The second question is do we really need to decompose step checking into
several stages. To answer the question, we design the following prompt to do
single-stage step checking.

> The following is a question and the first a few steps in its solution.
> Question: [Question]
> Solution: [Step 0, Step 1,…, Step i-1]
> Check the correctness of the next step: [Step i]
> Please consider the information it relies on and check step by step. Please
> end your response with your conclusion that starts with "Correct", "Wrong"
> or "Not Sure".

Figure 6 and Table 1 show though single-stage step checking is better than
global checking, it is still significantly worse than our multi-stage
SelfCheck. It indicates that the difficulty of checking a step is still too
much for the LLM, so we need to decompose step checking into a pipeline of
easier tasks. Figure 6

#### Regenerate then compare vs error check

Method | Accuracy (%)
---|---
Ours | 66.7%
Global | $55.0\%$
Single stage | $57.2\%$
Regen->check (0-shot) | $63.1\%$
Regen->check (1-shot) | $64.2\%$
Table 1: Mean checking accuracies for different variants of SelfCheck.

We also need to justify our choice to replace direct error checking with step
regeneration and comparison. To do so, we replace the regeneration stage and
comparison stage with a verification stage. We first compare with a zero-shot
version of the variant with the following prompt:

> Given the following information:
> Information 0: [Information $I_{0}$]
> Information 1: [Information $I_{1}$]
> …
> Step 0: [Step $S_{0}$]
> Step 1: [Step $S_{1}$]
> …
> Check the correctness of the next step [Step i]
> Please check for grounding errors, reasoning errors and calculation errors
> step by step. Please end your response with your conclusion that starts with
> "Correct", "Wrong" or "Not Sure".

To make the verification stage stronger, we also compare a one-shot version
with an example from Ling et al. (2023). Results in Figure 6 and Table 1 show
that even with a very detailed and instructive example, direct verification
still performs worse than our regeneration style checking, which supports our
argument in Section 2.

##  6 Conclusion

In this paper, we introduce SelfCheck, which is a step-by-step checking scheme
for multi-step reasoning. By our multi-stage structure and the regeneration
style checking, SelfCheck is effective in checking for errors, which, in turn,
greatly increase prediction accuracies by weighted voting. Since SelfCheck
works well in the zero-shot setting, we conclude that LLMs can check their own
outputs without external supervision.

## References

  * Amini et al. (2019) Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Koncel-Kedziorski, Yejin Choi, and Hannaneh Hajishirzi.  MathQA: Towards interpretable math word problem solving with operation-based formalisms.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pp. 2357–2367, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1245.  URL https://aclanthology.org/N19-1245.
  * (2) Anthropic.  Meet claude.  URL https://www.anthropic.com/product.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Bubeck et al. (2023) Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al.  Sparks of artificial general intelligence: Early experiments with gpt-4.  _arXiv preprint arXiv:2303.12712_ , 2023.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_ , 2022.
  * Cobbe et al. (2021) 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_ , 2021.
  * Fu et al. (2023) Yao Fu, Litu Ou, Mingyu Chen, Yuhao Wan, Hao Peng, and Tushar Khot.  Chain-of-thought hub: A continuous effort to measure large language models’ reasoning performance.  _arXiv preprint arXiv:2305.17306_ , 2023.
  * Google (2023) Google.  Palm 2 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * He et al. (2020) Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen.  Deberta: Decoding-enhanced bert with disentangled attention.  In _International Conference on Learning Representations_ , 2020.
  * Hendrycks et al. (2021) 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.  In _Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2)_ , 2021.
  * Li et al. (2022) Yifei Li, Zeqi Lin, Shizhuo Zhang, Qiang Fu, Bei Chen, Jian-Guang Lou, and Weizhu Chen.  On the advance of making language models better reasoners.  _arXiv preprint arXiv:2206.02336_ , 2022.
  * Lightman et al. (2023) 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.
  * Ling et al. (2023) Zhan Ling, Yunhao Fang, Xuanlin Li, Zhiao Huang, Mingu Lee, Roland Memisevic, and Hao Su.  Deductive verification of chain-of-thought reasoning.  _arXiv preprint arXiv:2306.03872_ , 2023.
  * Lyu et al. (2023) Qing Lyu, Shreya Havaldar, Adam Stein, Li Zhang, Delip Rao, Eric Wong, Marianna Apidianaki, and Chris Callison-Burch.  Faithful chain-of-thought reasoning.  _arXiv preprint arXiv:2301.13379_ , 2023.
  * OpenAI (2023) OpenAI.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Paul et al. (2023) Debjit Paul, Mete Ismayilzada, Maxime Peyrard, Beatriz Borges, Antoine Bosselut, Robert West, and Boi Faltings.  Refiner: Reasoning feedback on intermediate representations.  _arXiv preprint arXiv:2304.01904_ , 2023.
  * Peng et al. (2023) Baolin Peng, Michel Galley, Pengcheng He, Hao Cheng, Yujia Xie, Yu Hu, Qiuyuan Huang, Lars Liden, Zhou Yu, Weizhu Chen, et al.  Check your facts and try again: Improving large language models with external knowledge and automated feedback.  _arXiv preprint arXiv:2302.12813_ , 2023.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023.
  * Wang et al. (2022) Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V Le, Ed H Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou.  Self-consistency improves chain of thought reasoning in language models.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Wei et al. (2021) Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le.  Finetuned language models are zero-shot learners.  In _International Conference on Learning Representations_ , 2021.
  * Weng et al. (2022) Yixuan Weng, Minjun Zhu, Shizhu He, Kang Liu, and Jun Zhao.  Large language models are reasoners with self-verification.  _arXiv preprint arXiv:2212.09561_ , 2022.


marginparsep has been altered.
topmargin has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Return-Aligned Decision Transformer

Tsunehiko Tanaka 1 2  Kenshi Abe 2  Kaito Ariu 2  Tetsuro Morimura 2  Edgar
Simo-Serra 1

††footnotetext:  1Waseda University, Tokyo, Japan 2CyberAgent, Tokyo, Japan.
Correspondence to: Tsunehiko Tanaka<EMAIL_ADDRESS>

###### Abstract

Traditional approaches in offline reinforcement learning aim to learn the
optimal policy that maximizes the cumulative reward, also known as return.
However, as applications broaden, it becomes increasingly crucial to train
agents that not only maximize the returns, but align the actual return with a
specified target return, giving control over the agent’s performance. Decision
Transformer (DT) optimizes a policy that generates actions conditioned on the
target return through supervised learning and is equipped with a mechanism to
control the agent using the target return. Despite being designed to align the
actual return with the target return, we have empirically identified a
discrepancy between the actual return and the target return in DT. In this
paper, we propose Return-Aligned Decision Transformer (RADT), designed to
effectively align the actual return with the target return. Our model
decouples returns from the conventional input sequence, which typically
consists of returns, states, and actions, to enhance the relationships between
returns and states, as well as returns and actions. Extensive experiments show
that RADT reduces the discrepancies between the actual return and the target
return of DT-based methods.

##  1 Introduction

Offline reinforcement learning (RL) focuses on learning optimal policies using
trajectories from offline datasets Levine et al. (2020); Fujimoto & Gu (2021);
Jin et al. (2021); Xu et al. (2022). Many methods in offline RL traditionally
aim to learn the optimal policy that maximizes the cumulative reward, also
known as return. However, in various scenarios, it is crucial to train an
agent to align the cumulative reward, which we refer to as the actual return,
with a given target value, which we call target return. Through specifying the
actual return, it is possible to adjust the performance of game bots Shen et
al. (2021) or educational tools Singla et al. (2021) to match user skill
levels, or model the trajectories of non-expert agents in traffic simulation
analysis Nguyen et al. (2021). Despite this importance, existing approaches
exhibit discrepancies between the actual return and target return, which
significantly lower the quality of the agents. In this work, we aim to align
the actual return with the target return, enabling control of the agent’s
performance based on the target return.

Figure 1: Comparison between DT and the proposed RADT architecture. RADT
decouples the returns from states and actions in the input sequence,
explicitly modeling the relationships between returns and other modalities.

Figure 2: Absolute errors between the target return and the actual return in
both the MuJoCo and Atari domains. The discrepancies are normalized by the
return range between the top 5% and bottom 95% within the dataset. RADT
significantly reduces the discrepancies in both domains.

Decision Transformer (DT) Chen et al. (2021) optimizes a policy that generates
actions conditioned on the target return through supervised learning, and is
equipped with a mechanism to control the agent using the target return.
Specifically, this model takes a sequence comprising future desired returns,
past states, and actions as inputs, and outputs actions using a transformer
architecture Vaswani et al. (2017). The self-attention mechanism in the
transformer extracts significant features of the input sequence based on the
relative importance of each token to all the other tokens. DT integrates
returns into the input sequence to condition action generation. However, as
illustrated in Fig. 2, there are discrepancies between the actual returns and
the target returns. This error could be explained by the simultaneous handling
of returns with the other modalities such as states and actions. This
simultaneous processing could lead the model to overly focus on state or
action tokens, and it may become unable to align the actual return with the
target return appropriately.

In this paper, we propose a novel architecture, Return-Aligned Decision
Transformer (RADT), designed to align the actual return with the target
return, as shown in Fig. 2. The core idea is to use an architecture that
separates returns from the input sequence and explicitly models the
relationships between the returns and the other modalities. This architecture
prevents the reduction of the target return’s influence on the output actions
due to other modalities such as state and action. To achieve this, we employ
two techniques. The first is a unique cross-attention mechanism that focuses
on the relationship between a state-action sequence and a return sequence. The
second is adaptive layer normalization, which scales the state-action features
using parameters inferred solely from the return features. They are
specifically designed for seamless integration into the transformer block,
thus maintaining the simple supervised learning approach of DT. In our
experiments, RADT shows a significant reduction in the absolute error between
the actual return and the target return, decreasing it to 39.7% of DT’s error
in the MuJoCo Todorov et al. (2012) domain and 29.8% in the Atari Bellemare et
al. (2013) domain. We evaluate the key techniques of RADT through an ablation
study, demonstrating that each technique is effective for different tasks and
can complement each other.

Our contributions are summarized as follows:

  * •

We propose a novel approach for DT designed to align the actual return with a
given target return.

  * •

To realize this concept, we present Return-Aligned Decision Transformer, which
decouples the returns from states and actions in the input sequence,
explicitly modeling the relationships between returns and other modalities.

  * •

We introduce a unique cross-attention mechanism and adaptive layer
normalization to model the relationships between returns and other modalities.

  * •

Our experiments demonstrate that our method outperforms existing approaches in
aligning actual returns with target returns.

##  2 Preliminary

We assume a finite horizon Markov Decision Process (MDP) with horizon $T$ as
our environment, which can be described as
$\mathcal{M}=\langle\mathcal{S},\mathcal{A},\mu,P,R\rangle$, where
$\mathcal{S}$ represents the state space; $\mathcal{A}$ represents the action
space; $\mu\in\Delta(\mathcal{S})$ represents the initial state distribution;
$P:\mathcal{S}\times\mathcal{A}\to\Delta(\mathcal{S})$ represents the
transition probability distribution; and
$R:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ represents the reward function.
The environment begins from an initial state $s_{1}$ sampled from a fixed
distribution $\mu$. At each timestep $t\in[T]$, an agent takes an action
$a_{t}\in\mathcal{A}$ in response to the state $s_{t}\in\mathcal{S}$,
transitioning to the next state $s_{t+1}\in\mathcal{S}$ with the probability
distribution $P(\cdot|s_{t},a_{t})$. Concurrently, the agent receives a reward
$r_{t}=R(s_{t},a_{t})$.

Decision Transformer (DT) introduces the paradigm of transformers into the
context of offline reinforcement learning. At each timestep $t$ in the
inference, DT takes a sequence of desired returns, past states, and actions as
inputs, and outputs an action $a_{t}$. The input sequence of DT is represented
as

$\displaystyle\tau=(\hat{R}_{1},s_{1},a_{1},\hat{R}_{2},s_{2},a_{2},...,\hat{R}_{t},s_{t}),$
(1)

where $\hat{R}_{t}$ represents the returns computed over the remaining
steps111 Despite Eq. (1) taking inputs from $1$ to $t$, in practice, only the
last $K$ timesteps are processed.. It is calculated as
$\hat{R}_{t}=R^{\mathrm{target}}-\sum_{t^{\prime}=1}^{t-1}r_{t^{\prime}}$.
$R^{\mathrm{target}}$ is a given constant222$R^{\mathrm{target}}$ is the total
return of the trajectory in the dataset during training, and it is a given
constant during inference., which is the total desired return to be obtained
in an episode of length $T$. We refer to $R_{\mathrm{target}}$ as a target
return. Raw inputs, referred to as tokens, are individually projected into the
embedding dimension by separate learnable linear layers for return, state, and
action respectively, to generate token embeddings. The tokens are processed
using a Transformer-based GPT model Radford et al. (2018). The model is
trained using either cross-entropy or mean-squared error loss, calculated
between the predicted action and the ground truth from the offline datasets.

Figure 3: An overview of our proposed Return-Aligned Decision Transformer
(RADT). We split input tokens into the return sequence and state-action
sequence. RADT equips the transformer block with cross-attention and adaptive
layer normalization to highlight returns tokens. Figure 4: The detailed
structure of our cross-attention. In the cross-attention, the state-action
sequence is input as the query, and the return sequence serves as both key and
value. The query input $q$ and output $z$ of the multi-head attention are fed
into a linear layer to calculate the weight $\alpha$ for the output $z$.

The transformer is an architecture designed for processing sequential data,
including the attention mechanism, residual connection, and layer
normalization. The attention mechanism processes three distinct inputs: the
query, the key, and the value. This process involves weighting the value by
the normalized dot product of the query and the key. The $i$-th output token
of the attention mechanism is calculated as follows:

$\displaystyle z_{i}=\sum_{j=1}^{n}\mathrm{softmax}({\langle
q_{i},k_{\ell}\rangle}^{n}_{\ell=1})_{j}\cdot v_{j},$ (2)

DT uses self-attention, where query, key, and value are obtained by linearly
different transformations of the input sequence,

$\displaystyle q_{i}=\tau_{i}W_{i}^{Q},\ k_{i}=\tau_{i}W_{i}^{K},\
v_{i}=\tau_{i}W_{i}^{V},$ (3)

Moreover, it employs a causal mask to prohibit attention to subsequent tokens,
rendering tokens in future timesteps ineffective for action prediction.

Layer normalization standardizes the token features to stabilize learning.
Residual connection avoids gradient vanishing by adding the input and output
of attention layers or feed-forward layers. For further details on DT, refer
to the original paper Chen et al. (2021).

Our goal is to minimize the following absolute error between the target return
and the actual return for any given target return $R^{\mathrm{target}}$ with a
single model:

$\displaystyle\mathbb{E}\left[\left|R^{\mathrm{target}}-\sum_{t=1}^{T}r_{t}\right|\right].$
(4)

We refer to $\sum_{t=1}^{T}r_{t}$ as an actual return.

##  3 Return-Aligned Decision Transformer

In this section, we present the Return-Aligned Decision Transformer (RADT), a
model designed to align the actual return with the target return. While DT
constructs the input sequence using returns, states, and actions, we propose
to separate returns from the input sequence. RADT incorporates techniques to
explicitly model the relationships between the returns and the other
modalities, states and actions. The training process follows a supervised
learning paradigm of DT. We show the overview of RADT in Fig. 3.

###  3.1 Overview

We separately input the return sequence $\phi$ and the state-action sequence
$\tau^{-}$ into our model.

$\displaystyle\phi$ $\displaystyle=(\hat{R}_{1},\hat{R}_{2},...,\hat{R}_{T}),$
(5) $\displaystyle\tau^{-}$
$\displaystyle=(s_{1},a_{1},s_{2},a_{2},...,s_{T}).$ (6)

DT integrates returns, states, and actions into a single input sequence to
generate actions based on returns. However, this approach could diminish the
attention allocated to the return tokens, leading to discrepancies between the
actual return and the target return. This is because the self-attention
mechanism excessively focuses on modalities other than returns. To address
this, we split input tokens into the return and state-action sequences and
apply an architecture that focuses more on the returns. The details of our
architecture will be explained in the following section.

###  3.2 Architecture

We propose two techniques for RADT to explicitly model the relationships
between the returns and the other modalities: cross-attention and adaptive
layer normalization. Both techniques can be seamlessly integrated into
transformer blocks, and we find that employing them in tandem yields superior
performance. The architecture utilizing both techniques is illustrated on the
right side of Fig. 3.

Cross-Attention between Returns and States-Actions.  We apply cross-attention
between the state-action sequence $\tau^{-}$ and the return sequence $\phi$.
The detailed structure of our cross-attention is shown in Fig. 4. In our
cross-attention, the state-action sequence $\tau^{-}$ acts as the query, while
the return sequence $\phi$ acts as both the key and value,

$\displaystyle q_{i}=\tau^{-}_{i}W_{i}^{Q},\ k_{i}=\phi_{i}W_{i}^{K},\
v_{i}=\phi_{i}W_{i}^{V}.$ (7)

These query, key, and value are input into multi-head attention (Eq. 2), and
we obtain the output $z$. Note that we use a causal mask to ensure that tokens
in the state-action sequence $\tau^{-}$ cannot access future return tokens.

Following the powerful cross-attention-based technique for integrating two
different types of features Nguyen et al. (2022), we adaptively adjust the
scale for the output against the input of the multi-head attention. We
concatenate the multi-head attention input $q$ and output $z$ as a column
vector $[z_{i};q_{i}]\in\mathbb{R}^{2D}$ and obtain dimension-wise scaling
parameters $\alpha$ through a learnable affine projection. Using the scale
$\alpha$, we define residual connection as

$\displaystyle\alpha_{i}$ $\displaystyle=W[z_{i};q_{i}]+b,$ (8) $\displaystyle
q^{\mathrm{updated}}_{i}$ $\displaystyle=(1+\alpha_{i})\otimes z_{i}+q_{i},$
(9)

where $W\in\mathbb{R}^{D\times 2D}$ and $b\in\mathbb{R}^{D}$ are learnable
parameters, and $\otimes$ denotes the Hadamard product. The choice of
$1+\alpha_{i}$ allows the model to start with a baseline scaling of 1 (simple
addition) by zero-initialization of $W$ and $b$, resulting in the scale of
$z_{i}$ and $q_{i}$ being the same. This provides a stable starting point,
from which the model can learn to adaptively adjust the scaling. As the
training progresses, $\alpha_{i}$ is updated to refine the balance between
$z_{i}$ and $q_{i}$.

In the self-attention in DT, the input sequence $\tau$, which is a
concatenation of returns, states, and actions, is input as the query, key, and
value. This setup attempts to express all interrelations solely through
$W^{Q},W^{K},W^{V}$. However, the model gets distracted from such complicated
relationships, making it less responsive to the returns. We address this by
using cross-attention to focus on the relationships between returns and the
other modalities. Additionally, by introducing a scale $\alpha$ for the output
$z$ relative to the input query $q$, we can more flexibly represent the
relationship between the returns sequence $\phi$ and the state-action sequence
$\tau^{-}$.

Adaptive Layer Normalization. We next introduce adaptive layer normalization,
which adaptively changes depending on the returns. This involves replacing the
standard layer normalization layers in transformer blocks with adaptive layer
normalization layers.

$\displaystyle
q^{\prime}_{i}=(1+\gamma_{i})\otimes\mathrm{LayerNorm}(q_{i})+\beta_{i}$ (10)

where $q_{i},q^{\prime}_{i}$ denote the input and output of adaptive layer
normalization, respectively. $\gamma,\beta$ are dimension-wise scaling
parameters regressed by a linear layer from the return sequence $\phi$.
Through this regression, the model can condition the state-action sequence
$\tau^{-}$ with the return. Similar to the $1+\alpha_{i}$ in Eq. (8), by zero-
initializing the parameters of the linear layer, Eq. (10) can be considered
the same as the standard layer normalization at the beginning of training.

In DT, layer normalization is positioned before either the attention layer or
the feed-forward layer. This configuration creates pathways that bypass layer
normalization via a residual connection. We place adaptive layer normalization
after residual connection to ensure that all pathways undergo adaptive layer
normalization, as in the vanilla transformer Vaswani et al. (2017).

##  4 Experiments

Figure 5: Absolute Error Comparison of RADT and baselines in MuJoCo domain.
Each column represents a task. The x-axis is the target return. Target returns
are divided into seven equally spaced points based on the cumulative reward of
trajectories in the dataset, ranging from the bottom 5% to the top 5%. In this
graph, the bottom 5% is represented as 0, and the top 5% as 100. The y-axis
represents the absolute error between the actual return and the target return
obtained from our simulations. We report the mean and standard error over
three seeds.

Figure 6: Histogram of the cumulative reward of trajectories included in the
D4RL dataset. Each column represents a MuJoCo task. The x-axis is normalized
by the top 5% and bottom 5% of the dataset’s trajectory cumulative reward,
similar to Fig. 6. When viewed alongside Fig. 6, we observe that the absolute
errors of DT and DC are inversely proportional to the amount of data, RADT
consistently maintains a low absolute error regardless of the amount of data.
Figure 7: Absolute Error Comparison of RADT and baselines in the Atari domain.
The way to read these graphs is the same as in Fig. 6. We report the mean and
standard error over three seeds.

In this section, we conduct extensive experiments to evaluate the performance
of our RADT. First, we verify that RADT is effective in earning returns
consistent with various given target returns, compared to other baselines.
Next, we demonstrate through an ablation study that the two types of
architectural designs constituting RADT are effective individually, and that
using both types together further improves performance. Additionally, we show
that during training, RADT converges to the given target promptly.

###  4.1 Datasets

We evaluate RADT on continuous (MuJoCo Todorov et al. (2012)) and discrete
(Atari Bellemare et al. (2013)) control tasks in the same way as DT. MuJoCo is
characterized by a continuous action space with dense rewards. We use four gym
locomotion tasks from the widely-used D4RL Fu et al. (2020) dataset: Ant,
Hopper, HalfCheetah, and Walker. In this experiment, we utilize the v2 medium-
replay dataset. The medium-replay comprises a replay buffer of policies that
achieve approximately $1/3$ of expert performance. Atari involves high-
dimensional visual observations and requires capturing long-term context to
handle the delay between actions and the resulting rewards. We use four tasks:
Breakout, Pong, Qbert, and Seaquest. Similar to DT, we use 1% of all samples
in the DQN-replay datasets as per Agarwal et al. (2020) for training.

###  4.2 Baselines and Settings

To evaluate our proposed model architecture, we use return-conditioned DT
models with different architectures, specifically DT Chen et al. (2021),
StARformer Shang et al. (2022), and Decision ConvFormer (DC) Kim et al.
(2024), as baselines. We use the official PyTorch implementations for
baselines. Although StARformer uses step-by-step rewards instead of returns,
in our experiments, we employ return-conditioning using returns. This
modification is stated in the original paper Shang et al. (2022) to have
minimal impact on performance. The hyperparameters for each method are set
according to the defaults specified in their respective papers or open-source
codebases. For visual observations in the Atari domain, RADT and DC use the
same CNN encoder as DT. StARformer, in addition to the CNN encoder, also
incorporates more powerful patch-based embeddings like Vision Transformer
Dosovitskiy et al. (2021).

In MuJoCo, for each method, we train three instances with different seeds, and
each instance runs 20 episodes for each target return. In Atari, for each
method, we train three instances with different seeds, and each instance runs
10 episodes for each target return. Target returns are set by first
identifying the range of cumulative reward in trajectories in the training
dataset, specifically from the bottom 5% to the top 5%. This identified range
is then equally divided into seven intervals, not based on percentiles, but by
simply dividing the range into seven equal parts. Each of these parts
represents a target return.

Table 1: Results of the ablation study on the proposed techniques proposed in
Sec.3.3. Each value indicates the sum of absolute error between the actual
return and the target return, relative to the target returns in the main
results. We report the mean and standard error over 3 seeds, and normalize
such that the average and variance of the results when using both techniques
(RADT) are 1.0 respectively. CA represents cross-attention, and AdaLN stands
for adaptive layer normalization.
Approach | CA | AdaLN | Ant | HalfCheetah | Hopper | Walker
---|---|---|---|---|---|---
DT |  |  | 2.69$\pm$2.07 | 2.66$\pm$2.07 | 2.19$\pm$6.49 | 2.39$\pm$2.60
RADT w/o AdaLN | ✓ |  | 1.29$\pm$0.38 | 1.51$\pm$1.19 | 1.20$\pm$1.55 | 1.24$\pm$2.66
RADT w/o CA |  | ✓ | 1.25$\pm$0.33 | 1.17$\pm$1.77 | 1.61$\pm$1.72 | 1.26$\pm$2.48
RADT Full | ✓ | ✓ | 1.00$\pm$1.00 | 1.00$\pm$1.00 | 1.00$\pm$1.00 | 1.00$\pm$1.00
Figure 8: Transitions of the absolute errors between the actual and target
returns during training in HalfCheetah, comparing the absolute errors of DT,
StARformer, and DC. The mean and standard error are reported across 3 seeds.

###  4.3 Main Results

The main results are presented in Fig. 6 for the MuJoCo domain and Fig. 7 for
the Atari domain. These figures plot the absolute error between the actual
return and the target return, where lower is better. Additionally, the target
returns are represented with the top 5% values as 100 and the bottom 5% values
as 0.

MuJoCo Domain. In the MuJoCo domain, as shown in Fig. 6, RADT outperforms the
baseline in all tasks. In HalfCheetah, the absolute error of DT becomes
smaller as the target return gets closer to 0 or 100, while in other tasks,
the error decreases as the target return gets closer to 0. Additionally, we
observe that DC exhibits a similar trend in absolute error to DT across all
four tasks. These trends in absolute error are inversely proportional to the
cumulative reward of trajectories in the dataset, as shown in Fig. 6. We
believe that DT and DC are susceptible to the distribution of data. In
contrast, RADT consistently shows smaller absolute errors across various
target returns compared to DT and DC, indicating its robustness to data
distribution.

It is noteworthy that in StARformer, the absolute error decreases as it gets
close to 66.7 in Hopper and 100 in HalfCheetah. This is because, regardless of
the specified target return, the actual return tends to converge to 66.7
(Hopper) or 100 (HalfCheetah), as shown in Fig. 9 in Appendix A.1. Shang et
al. (2022) report that it performs equally well with or without using returns
or step-by-step rewards. We believe that StARformer generates actions
independently of returns or step-by-step rewards, suggesting that it is
challenging to control it with various target returns. Conversely, RADT does
not exhibit the bias seen in StARformer and responds more accurately to the
target returns.

Atari Domain. The Atari domain is challenging due to the sparsity of rewards,
making the control of obtained rewards more difficult than in MuJoCo. Despite
this, RADT outperforms the baselines in most tasks, as shown in Fig. 7.
Moreover, given that observations in the Atari domain are images, enhancements
to the vision encoder offer considerable advantages. However, when compared to
StARformer, which utilizes powerful patch-based embeddings such as ViT, RADT
shows superior performance in most tasks and achieves competitive results in
others. DC reports excellent results for Pong, demonstrating that convolution
is superior to self-attention in limited tasks. Meanwhile, our proposed method
performs best in the other tasks, suggesting that RADT would be better suited
for a wide range of scenarios.

###  4.4 Ablation Study

Architecture Techniques. We conduct an ablation study for the core techniques
of RADT, as proposed in Sec. 3.2, on the MuJoCo domain. The results are
summarized in Tab. 1. CA represents cross-attention, and AdaLN stands for
adaptive layer normalization. Each value is the sum of the absolute errors
between the actual return and the target return across the multiple target
returns. These values are then normalized so that the average and standard
error of the results when using both techniques (RADT) are 1.0 respectively.
Table 1 reports that introducing either cross-attention or adaptive layer
normalization individually proves to be effective. The cross-attention
performs better in Hopper and Walker, while adaptive layer normalization
excels in Ant and HalfCheetah. This suggests that each technique has its
unique areas of expertise. The combination of both techniques results in the
smallest error across all tasks, implying that cross-attention and adaptive
layer normalization effectively complement each other.

Training Convergence. We present in Fig. 8 the curves of the absolute errors
between the actual return and the target return as training proceeds, along
with the absolute errors of DT, StARformer, and DC as baselines. For all
target returns, we observe that RADT converges up to 20k iterations. This
suggests that RADT has an advantage in terms of the speed and stability of
training compared to other baselines.

##  5 Related Work

###  5.1 RTG-Conditioned Offline RL

Recent studies have focused on formulating offline reinforcement learning (RL)
as a problem of predicting action sequences that are conditioned by goals and
rewards Chen et al. (2021); Emmons et al. (2022); David et al. (2023);
Schmidhuber (2019); Srivastava et al. (2019). This approach differs from the
popular value-based methods Kumar et al. (2020); Kostrikov et al. (2022) by
modeling the relationship between rewards and actions through supervised
learning. Decision Transformer (DT) Chen et al. (2021) introduces the concept
of desired future returns and improves performance by training the transformer
Vaswani et al. (2017) as a return-conditioned policy. There have been various
advancements in offline RL based on the DT framework, which will be elaborated
in the following section.

###  5.2 Transformers for RL

Some efforts focus on refining the transformer architecture for offline RL.
StARformer Shang et al. (2022) introduces two transformer architectures, one
aggregates information at each step, and the other aggregates information
across the entire sequence. The image encoding process is improved by dividing
the observation images into patches and feeding them into the transformer to
enhance step information, similar to Vision Transformer Dosovitskiy et al.
(2021). Decision ConvFormer Kim et al. (2024) replaces attention with
convolution to capture the inherent local dependence pattern of MDP. While
these architectures preserve the input sequence structure of the transformer,
comprising returns, states, and actions, they do not directly tackle the
challenge of diminishing the influence of returns on the decision-making
process. In contrast, our research specifically aims to align the actual
return with the target return.

There is growing interest in methods of manipulating data to train DT to solve
the trajectory stitching problem in offline RL. Stitching refers to the
ability to combine sub-optimal trajectories to create an optimal one. For
example, Q-learning Decision Transformer (QDT) Yamagata et al. (2023) and
Advantage Conditioned Transformer (ACT) Gao et al. (2023) re-label the data
trajectories using return-to-go or advantage obtained by dynamic programming.
Elastic Decision Transformer (EDT) Wu et al. (2023) adjusts the length of
trajectories maintained in DT to facilitate the stitching of trajectories
during testing. Liu & Abbeel (2023) prepare multiple trajectories sorted in
ascending order by total rewards and re-labels them with the highest total
reward in the trajectories before training DT. These efforts enhance the
stitching capability, demonstrating that agents trained in offline RL can
achieve better performance. While our focus is on model improvement and thus
our scope differs from these efforts, we anticipate that combining these
efforts and our model could improve the performance of both.

Several training methods have been proposed that emphasize the performance of
DT agents when transitioning from offline training to online activity. Online
Decision Transformer Zheng et al. (2022) fine-tunes policies pre-trained on
offline datasets through online interactions with the environment. Janner et
al. (2021) add beam search to enable the agent to explore better actions. Wang
et al. (2023) introduce a trained value function into DT, bridging the gap
between the deterministic nature of return-conditioned supervised learning and
the probabilistic characteristics of value-based methods. While these online
enhancements diverge from the primary focus of this study, we believe that a
combination of this study with these enhancements can potentially open up
further possibilities.

##  6 Discussion

Limitations. As a direction for future improvements, one may enhance the
performance of RADT in maximizing actual return. The comparison between RADT
and other baselines in terms of maximizing the actual return is summarized in
Tab. 2. We consider five baselines: CQL Kumar et al. (2020), IQL Kostrikov et
al. (2022), DT Chen et al. (2021), DC Kim et al. (2024), and ODT Zheng et al.
(2022). RADT shows results that are either surpassing or competitive with DT.
However, there is still a performance gap compared to value-based methods like
CQL and IQL, and DT’s variants such as DC and ODT. To bridge this gap, our
proposed method can easily be combined with other techniques because our
proposed method is fully encapsulated within a transformer block. For
instance, there is potential to enhance performance further by fine-tuning
RADT with the online approach ODT. While this direction is interesting, it’s
important to note that RADT’s primary aim is for the actual return to align
with the target return, and maximizing the actual return is beyond the main
scope of this paper.

Table 2: Performance comparison in the actual return maximization in the
MuJoCo domain. We cite the results for CQL, DT, DC, and ODT from their
reported scores. The results for IQL are cited from Kostrikov et al. (2022).
We report the average and standard error across 3 seeds from our simulation
results.
Method | HalfCheetah | Hopper | Walker
---|---|---|---
CQL | 45.5 | 95.0 | 77.2
IQL | 44.2 | 94.7 | 73.9
DC | 41.3 | 94.2 | 76.6
ODT | 41.5 | 91.9 | 81.0
DT | 36.6 | 82.7 | 66.6
RADT | 39.2$\pm$0.7 | 88.0$\pm$1.8 | 66.2$\pm$2.1

Conclusion. We have explored the development of advanced, controllable agents
in offline RL based on DT. Our proposed Return-Aligned Decision Transformer
(RADT) separates the return sequence from the state-action sequence and can
significantly improve the alignment of actual returns with the target return.
The architectural innovations of RADT are a distinctive cross-attention
mechanism and adaptive layer normalization for returns, designed to be easily
integrated into the transformer block. In both MuJoCo and Atari domains, RADT
demonstrated superior capability in aligning the actual return with the target
return compared to existing DT-based methods. We believe that this capability
allows us to obtain agents at various skill levels, and it broadens the range
of applications for DT, such as in traffic analysis and as tools for games and
education.

## References

  * Agarwal et al. (2020) Agarwal, R., Schuurmans, D., and Norouzi, M.  An optimistic perspective on offline reinforcement learning.  In _Proc. of ICML_ , 2020.
  * Bellemare et al. (2013) Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47(1):253–279, 2013.
  * Chen et al. (2021) Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I.  Decision transformer: Reinforcement learning via sequence modeling.  In _Proc. of NeurIPS_ , 2021.
  * David et al. (2023) David, S. B., Zimerman, I., Nachmani, E., and Wolf, L.  Decision S4: Efficient sequence-based RL via state spaces layers.  In _Proc. of ICLR_ , 2023.
  * Dosovitskiy et al. (2021) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _Proc. of ICLR_ , 2021.
  * Emmons et al. (2022) Emmons, S., Eysenbach, B., Kostrikov, I., and Levine, S.  Rvs: What is essential for offline RL via supervised learning?  In _Proc. of ICLR_ , 2022.
  * Fu et al. (2020) Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Fujimoto & Gu (2021) Fujimoto, S. and Gu, S.  A minimalist approach to offline reinforcement learning.  In _Proc. of NeurIPS_ , 2021.
  * Gao et al. (2023) Gao, C., Wu, C., Cao, M., Kong, R., Zhang, Z., and Yu, Y.  Act: Empowering decision transformer with dynamic programming via advantage conditioning.  _arXiv preprint arXiv:2309.05915_ , 2023.
  * Janner et al. (2021) Janner, M., Li, Q., and Levine, S.  Offline reinforcement learning as one big sequence modeling problem.  In _Proc. of NeurIPS_ , 2021.
  * Jin et al. (2021) Jin, Y., Yang, Z., and Wang, Z.  Is pessimism provably efficient for offline RL?  In _Proc. of ICML_ , 2021.
  * Kim et al. (2024) Kim, J., Lee, S., Kim, W., and Sung, Y.  Decision ConvFormer: Local filtering in MetaFormer is sufficient for decision making.  In _Proc. of ICLR_ , 2024.
  * Kostrikov et al. (2022) Kostrikov, I., Nair, A., and Levine, S.  Offline reinforcement learning with implicit Q-learning.  In _Proc. of ICLR_ , 2022.
  * Kumar et al. (2020) Kumar, A., Zhou, A., Tucker, G., and Levine, S.  Conservative Q-learning for offline reinforcement learning.  In _Proc. of NeurIPS_ , 2020.
  * Levine et al. (2020) Levine, S., Kumar, A., Tucker, G., and Fu, J.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Liu & Abbeel (2023) Liu, H. and Abbeel, P.  Emergent agentic transformer from chain of hindsight experience.  In _Proc. of ICML_ , 2023.
  * Nguyen et al. (2021) Nguyen, J., Powers, S. T., Urquhart, N., Farrenkopf, T., and Guckert, M.  An overview of agent-based traffic simulators.  _Transportation Research Interdisciplinary Perspectives_ , 12:100486, 2021.
  * Nguyen et al. (2022) Nguyen, V.-Q., Suganuma, M., and Okatani, T.  Grit: Faster and better image captioning transformer using dual visual features.  In _Proc. of ECCV_ , 2022.
  * Radford et al. (2018) Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al.  Improving language understanding by generative pre-training.  2018\.
  * Schmidhuber (2019) Schmidhuber, J.  Reinforcement learning upside down: Don’t predict rewards–just map them to actions.  _arXiv preprint arXiv:1912.02875_ , 2019.
  * Shang et al. (2022) Shang, J., Kahatapitiya, K., Li, X., and Ryoo, M. S.  StARformer: Transformer with state-action-reward representations for visual reinforcement learning.  In _Proc. of ECCV_ , 2022.
  * Shen et al. (2021) Shen, R., Zheng, Y., Hao, J., Meng, Z., Chen, Y., Fan, C., and Liu, Y.  Generating behavior-diverse game ais with evolutionary multi-objective deep reinforcement learning.  In _Proc. of IJCAI_ , 2021.
  * Singla et al. (2021) Singla, A., Rafferty, A. N., Radanovic, G., and Heffernan, N. T.  Reinforcement learning for education: Opportunities and challenges.  _arXiv preprint arXiv:2107.08828_ , 2021.
  * Srivastava et al. (2019) Srivastava, R. K., Shyam, P., Mutz, F., Jaśkowski, W., and Schmidhuber, J.  Training agents using upside-down reinforcement learning.  _arXiv preprint arXiv:1912.02877_ , 2019.
  * Todorov et al. (2012) Todorov, E., Erez, T., and Tassa, Y.  MuJoCo: A physics engine for model-based control.  In _Proc. of IROS_ , 2012.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  In _Proc. of NeurIPS_ , 2017.
  * Wang et al. (2023) Wang, Y., Yang, C., Wen, Y., Liu, Y., and Qiao, Y.  Critic-guided decision transformer for offline reinforcement learning.  _arXiv preprint arXiv:2312.13716_ , 2023.
  * Wu et al. (2023) Wu, Y.-H., Wang, X., and Hamaya, M.  Elastic decision transformer.  In _Proc. of NeurIPS_ , 2023.
  * Xu et al. (2022) Xu, H., Jiang, L., Li, J., and Zhan, X.  A policy-guided imitation approach for offline reinforcement learning.  In _Proc. of NeurIPS_ , 2022.
  * Yamagata et al. (2023) Yamagata, T., Khalil, A., and Santos-Rodriguez, R.  Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline RL.  In _Proc. of ICML_ , 2023.
  * Zheng et al. (2022) Zheng, Q., Zhang, A., and Grover, A.  Online decision transformer.  In _Proc. of ICML_ , 2022.

##  Appendix A Additional Results

###  A.1 Further Visualizations of Main Results

Concerning the results of Fig. 6 and Fig. 7 presented in Sec. 4.3, we
illustrate the comparison of performance for each target return in Fig. 9. The
black dotted line represents $y=x$, indicating that the actual return matches
the target return perfectly. The closer to the black dotted line, the better
the result. In the MuJoCo domain, it is shown that there is a correlation
between the target return and the actual return for all methods. In contrast,
in the Atari domain, especially in the Qbert task, the baseline shows a lack
of correlation with the target return. Furthermore, there tends to be a
performance that exceeds the black dotted line, suggesting that aligning with
the target return is more difficult compared to the MuJoCo domain.

(a) MuJoCo.

(b) Atari.

Figure 9: Comparisons of actual returns per target return in the experiments
of Fig. 6 and Fig. 7. Each column represents a task. The x-axis represents the
target return, and the y-axis represents the actual return. Target returns are
set in the same way as Fig. 6 and Fig. 7. The black dotted line represents
$y=x$, indicating that the actual return matches the target return perfectly.
We report the mean and standard error over three seeds.

###  A.2 Scaling Parameters in Cross-Attention

We conduct experiments to evaluate the scaling parameter $\alpha$, included in
the cross-attention introduced in Sec. 3.2. The results are shown in Tab. 3.
In these experiments, we did not employ adaptive layer normalization in either
scenario, whether with or without $\alpha$ We observed that the introduction
of cross-attention without scaling improves the results for some tasks
compared to DT. Furthermore, the inclusion of the scaling parameter $\alpha$
significantly enhance performance, surpassing DT in all tasks.

Table 3: Experimental results evaluating the scaling parameters $\alpha$ in
the MuJoCo domain. CA represents DT with our cross-attention, and we prepare
versions with and without the scaling parameters. The results of DT are also
included.
Method | Scaling | Ant | HalfCheetah | Hopper | Walker
---|---|---|---|---|---
DT |  | 2.69$\pm$2.07 | 2.66$\pm$2.07 | 2.19$\pm$6.49 | 2.39$\pm$2.60
CA |  | 1.97$\pm$0.13 | 2.74$\pm$0.35 | 2.15$\pm$0.10 | 2.90$\pm$0.21
CA | ✓ | 1.29$\pm$0.38 | 1.51$\pm$1.19 | 1.20$\pm$1.55 | 1.24$\pm$2.66

##  Appendix B Experimental Details

###  B.1 Comparison of discrepancies

The discrepancies in Fig. 2 are calculated as the sum of the absolute errors
between the actual return and target return across multiple tasks. The
absolute error is normalized by the difference between the top 5% and bottom
95% of returns within the dataset.

###  B.2 Hyperparameter Settings

The full list of hyperparameters for RADT can be found in Tab. 4 and Tab. 5.
We use the model code for DT, StARformer, and DC from the following sources.
DT: https://github.com/kzl/decision-transformer. StARformer:
https://github.com/elicassion/StARformer. DC:
https://openreview.net/forum?id=af2c8EaKl8.

Table 4: Hyperparameters settings of RADT in the MuJoCo domain.
Hyperparameter | Value
---|---
Number of layers | $3$
Number of attention heads | $1$
Embedding dimension | $128$
Nonlinearity function | GELU, transformer
| SiLU, adaptive layer normalization
Context length K | $20$
Dropout | $0.1$
Learning rate | $10^{-4}$
Grad norm clip | $0.25$
Weight decay | $10^{-4}$
Learning rate decay | Linear warmup for first $10^{5}$ training steps
Position Encoding | Sinusoidal Position Encoding
Table 5: Hyperparameters settings of RADT in the MuJoCo domain.
Hyperparameter | Value
---|---
Number of layers | $6$
Number of attention heads | $8$
Embedding dimension | $128$
Batch size |  $512$ Pong
|  $128$ Breakout, Qbert, Seaquest
Nonlinearity | ReLU encoder
| GELU transformer
| SiLU adaptive layer normalization
Encoder channels | $32,64,64$
Encoder filter size | $8\times 8,4\times 4,3\times 3$
Encoder strides | $4,2,1$
Max epochs | $5$
Dropout | $0.1$
Learning rate | $6\times 10^{-4}$
Adam betas | $(0.9,0.95)$
Grad norm clip | $0.1$
Weight decay | $0.1$
Learning rate decay | Linear warmup and cosine decay
Warmup tokens | $512*20$
Final tokens | $2*500000*K$
Position Encoding | Sinusoidal Position Encoding