Return-Aligned Decision Transformer
Authors: Tsunehiko Tanaka
Kenshi Abe
Kaito Ariu
Tetsuro Morimura
Edgar Simo-Serra
Venue: Accepted by TMLR ()

Abstract: Traditional approaches in offline reinforcement learning aim to learn the optimal policy that maximizes the cumulative reward, also known as return. It is increasingly important to adjust the performance of AI agents to meet human requirements, for example, in applications like video games and education tools. Decision Transformer (DT) optimizes a policy that generates actions conditioned on the target return through supervised learning and includes a mechanism to control the agent's performance using the target return. However, the action generation is hardly influenced by the target return because DT’s self-attention allocates scarce attention scores to the return tokens. In this paper, we propose Return-Aligned Decision Transformer (RADT), designed to more effectively align the actual return with the target return. RADT leverages features extracted by paying attention solely to the return, enabling action generation to consistently depend on the target return. Extensive experiments show that RADT significantly reduces the discrepancies between the actual return and the target return compared to DT-based methods.

Published in Transactions on Machine Learning Research (06/2025)
Return-Aligned Decision Transformer
Tsunchiko Tanaka Woseda University
tsunchiko@fuji wasedo.ip
Kenshi Abe CyberAgent
abekenshi1224 (atl//2222mmmlll. @gmaill com
Kaito Ariu Ageer
Tetsuro Morimura AyberAgent
morimara frriium__ldduurd@@geeenggg.....
Edgar Simo-Scrra Waseda University
cssChwascda..p
Reviewed on OpenReview: https: /ttp://pppp openreview. net/ forum? id= 17t2c7W8h1
Abstract
Traditional approaches offline reinforcement learning aim to learn the optimal policy that maximizes the cumulative reward, also known return. It. is increasingly important to adjust the performance of AT agents to meet human requirements, for example, in applica- lions like video games and education tools. Decision Transformer (DT) optimizes a policy that generates actions conditioned on the target return through supervised learning and includes mechanism to control the agent's performance using the target return. However, the action generation hardly influenced by the target return because DT's self-attention allocates scarce altention scores to the return tokens. In this paper, we propose Return- Aligned Decision Transformer (RADT), designed more effectively align the actual return with the target return. RADT leverages features extracted by paying attention solely to the return, enabling action generation to consistently depend on the larget return. Extensive experiments show that RADT significantly reduces the discrepancies between the actual return and the target return compared to DTT-based methods.
Introduction
Offline reinforcement learning (RL) focuses on learning policies from trajectories collected in offline datasets {Levine et al.] 22020] Fujimoto & Gul 2021 Yu et al. 2021 Jin et al. 2021 Xu et al. 2022) While many offline RL methods aim to optimize policies for maximum cumulative rewards (returns), they typically produce single policy fixed to a specific performance level. This rigidity problematic in scenarios requiring agents with varying skill levels, as forces developers to itcratively adjust reward functions and retrain models for each desired performance. Obtaining lower-performing policies from intermediate model checkpoints during training is also unsuitable for flexible performance adjustments. This is because accu- rately assossing their performance requires online evaluation, which is often difficult to conduct in practice. These limitations make it difficult to tailor policies to specific performance requirements.
To overcome these limitations, we propose training a single model that can achieve any desired target return. By simply adjusting the target refurn parameter, one can obtain agents spanning wide range of performance levels. This capability streamlines the policy development process in real-world applications, where heterogeneous performance often required:

---

Published in Transactions on Machine Learning Research (06/2025)
RADT (ours)
Return
state
action
Return
19
10
0.00
0.09
0:18
Figure 1: Comparison of attention scores between DT and RADT trained on the MuJoCo -nt-medium-rrelay dataset. The figure shows the attention scores assigned to input sequences by the first self-attention layer in DT and the first SeqRA in RADT. Each value represents the episode average of the attention scores. DT assigns minimal attention to return-to-go tokens, while RADT assigns attention exclusively to return-to-go tokens, thereby improving return alignment in action generation.
Example (Creating AI opponents in video games and oducational tools). Consider implementing AT opponents using RL in vidco games (Yannakakis & Togelius 2018) and educational tools (Singla ct al., 2021) development. Developers implement AT opponents with varied skill levels to ensure players, from beginners to experts, are appropriately challenged (Shen et al. 2021 Wu et al.) 2023a). Our proposed approach allows developers to interactively adjust the target return to produce a diverse set of AI opponents with different skill levels. This significantly reduces development overhead, enhancing game quality by providing tailored experiences for each player.
Example 2 (Human motion generation). Human motion generation in robotics and animation Zhu et al. 2024) often uses RL to produce natural and controllable human pose sequences (Jiang et al. 2024 Peng et al., 2022). Our approach enhances controllability by enabling a single policy to adjust its behavior based the target return. For example, using forward velocity as reward function, single policy model can generate walking motions for low target returns and rumning motions for high target returns. This flexibility improves the efficiency and versatility of motion synthesis for various downstream tasks (Rempe et al. 2023
Xie et al. 2024
For these applications, Decision Transformer (DT) (Chen et al. 2021 is a promising method that enables control over an agent's performance via a target return. DT optimizes a policy through supervised learning to generate actions conditioned on the target return. Specifically, it takes a sequence comprising desired future returns (also known as return-to-go), past states, and actions as inputs, and outputs actions using transformer architecture (Vaswani et al., 2017). In the self-attention mechanism of the transformer, each token selectively incorporates features from other tokens based on their relative importance, also known as the attention score. DT leverages the self-attention mechanism to propagate return-to-go tokens across the input sequence, enabling the conditioning of action generation on the target return. However, despite this conditioning, DT often obtains an actual return that diverges significantly from the target return. Our analysis, as ratest in Fig. [ reveals that the self-allention mechanism assigns only minimal attention to the return-to-go tokens, causing the return-to-go information to be nearly lost as it propagates through the network. This suggests that the return-to-go information nearly vanishes after passing through DT's self-attention, and leads DT to generate actions that are largely independent of the target return.
In this paper, we propose the Return-Aligned Decision Transformer (RADT), a novel architecture designed to align the actual return with the target return. Our key idea is to separate the return-to-go sequence from the state-action sequence so that the return-to-go sequence more directly influences action generation. To achieve this, we adopt two complementary design strategics: Sequence Return Aligner (SeqRA), described in Sec. and Stepwise Return Aligner (StepRA), discussed in Sec. 4.2 SeqRA processes return-to-go tokens from multiple past timesteps to capture long-term dependencies. In contrast, StepRA links each state action token directly to its corresponding return-to-go token at the same timestep, thereby capturing their stepwise relationship. By integrating these two approaches, RADT effectively leverages return-to-go information throughout the deeision-making. In our experiments, RADT siguificantly reduces the absolute error between the actual and target returns, achieving reductions of 54.9% and 34.4% compared to DT in the MuJoCo Todorov et al.] 2012) and Atari (Bellemare et al.] 2013) domains, respectively. Ablation studies demonstrate that each strategy independently contributes to improving return alignment, and their combination further enhances the model's ability to match the actual return to the target return.
2

---

Published in Transactions on Machine Learning Research (06/2025)
In summary, our contributions are follows:
1. We introduce RADT, a novel offlinc RL approach designed to align the actual return with the target return.
2. RADT employs unique architectural design that treats the return-to-go sequences distinctly from state-action sequences, enabling the return-to-go information to guide action generation effectively.
3. We present empirical evidence that RADT surpasses existing DTT-based models in return alignment.
2 Preliminary
We assumc finite horizon Markov Decision Process (MDP) with horizon T our environment, which can be described as M = (S, A, μt, P, R), where S represents the state space; A represents the action space; I E (SS) represents the initial state distribution; S x A + ASS) represents the transition probability distribution; and R : S x + represents the reward function. The environment begins from an initial state $1 sampled from fixed distribution μ.. At each timestep in response to the state 8t S, transitioning to the next state 8t+1 E S with the probability distribution P(+1s,,;;) Concurrently, the agent receives reward R(s,, a.).
[T], an agent takes an action a E A
introduces the paradigm of transformers in the context of offline reinforcement learning. We consider a constant Rtaret, which represents the total desirable return obtained throughout an episode of length T. We refer to Rtarget as the target return. At each timestep during inference, the desirable return to be obtained the remaining steps calculated as follows:
Decision Transformer (DT) (Chen et al. 2021
R = Rtargot
Σ P=1
(1)
We refer to R as return-to-go. DT takes sequence of return-to-go, past states, and actions as inputs, and outputs an action az- The input sequence of DTH is represented as
T= (fi, (1,,,1, 01, R2z 82, .2, Ri, 8t).
(2)
Raw inpuls, referred to as tokens, are individually projected into the embedding dimension by separate learn- able linear layers for returns-to-gr state, and action respectively, to generate token embeddings. Note that from this point onwards, we will denote tokens as R4 8;, a;- The tokens are processed using a transformer- based GPT model [Radford et al. 2018). The processed token 8t input into the prediction head to predict the action d₁- The model trained using either cross-entropy mean-squared error loss, calculated bet the predicted action and the ground truth from the offline datasets
The transformer (Vaswani et al. 2017) 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 weight also known as the attention score, and is caleulated follows:
0gj = softmax(((aa... k))
(3)
where Oij 0, Vj denotes a causal mask and n denotes the input length. The causal mask prohibits attention to subsequent tokens, rendering tokens in future timesteps ineffective for action prediction. The i-th output token of the attention mechanism is calculated as follows:
Z4
Oig Cg where Σ aij = and ai 2 0.
(4)
1=0
I For practicality, only the last. K timesteps are processed, rather than considering the full inputs. 2 ptarget is the total return the trajectory in the dataset during training.
3

---

Published in Transactions on Machine Learning Research (06/2025)
DT uses self-attention, where query, key, and value are obtained by linearly different transformations of the input sequence,
4i = T,W,, k = TWWK v; T,
(5)
Layer nonmalization standardizes the loken features stabilize learning. Residual connection avoids gradient vanishing by adding the imput and output of attention layers or beed-forward layers. For further details DT, refer to the original paper (Chen et al. 2021
DT conditions action gemeration by employing the sell-allention mechanism to disseminate relurn-to-go information throughout the input sequence. Despite its design, DT cannot align the actual return with the target return. To address this challenge, our goal is to minimize the following absolute error between the target return Rtarget and the actual return Σ(=1 rt using a single model:
E
ptaget
r1
(6)
3 Return Alignment Difficulty in DT
DT
RADT
Target Return
3160.5
3000
2619.9
1538.8
2000
998.2
457.7
Return-to-go
-82.9
1000
0
Overshooting target return
Reducing excess beyond target return
-1000
200
400
600
800
1000
200
400
600
800
1000
Timestep
Timestep
Figure 2: Comparison of return-to-go transitions with varying target returns. We evaluate DT and RADT on the ant-medium-replay dataset of MuJoCo. For each target return, we conduct 100 episodes, calculating the average and standard error of the target return at every step. The return-to-go of zero mcans that the actual return exactly matches the larget return.
In order understand the behavior of DT, we analyze how DT trained on the ant-mediun-replay dataset of MuJoCo responds the return-to-go throughout episode. Ideally, the return-to-go should be zero at the end of the episode, which means that the actual return perfectly matches the target return. Figure 2 plots the changes in return-to-go input to the model during an episode. As shown in Fig. 2 (left), DT reaches values of return-to-go that fall significantly below zero for all target returns, and obtains actual returns that differ greatly from the target returns.
We hypothesize that the problem lies the architoctural design of DT. Specifically, DT conditions action predictions on target returns through return-to-go lokens within the input sequence. These tokens compete for attention allocation in DT's self-attention mechanism, shown in Eq. 4 When attention is predom- inantly given to state or action tokens, less attention is allocated to return-to-go tokens, weakening the influence of target returns on action prediction.
This hypothesis is supported by analyzing attention scores from DT's first self-attention layer, as shown in Fig. (left). We observe that attention is strongly biased toward state tokens, with minimal attention return-to-go tokens. This tendency could arise because the training data (Fu et al. 2020) includes samples generated by Markov policies that do not take return input, causing DT to disregard return.
4

---

Published in Transactions on Machine Learning Research (06/2025)
4 Return-Aligned Decision Transformer
DT
RADT
--
Q--1
StepRA
Reed-Forwatt
Feed Forward
StepRA
DT Block
RADT Block
Self-Attention
SeqRA
StepRA
R
St-1
ar-1
R₁
St.
St-1
Q-11
St
Sell-Attention
R--1
Separate Return-to-ge tokens from State-Action Sequence
State-Action Tokens
Return-to-go Tokens
Figure 3: Comparison between DT and the proposed RADT architecture. DT processes a combined sequence relurns-to-go, states, and actions as input. In contrast, RADT separates the return-Lo-go from the state-action sequence and applies two Return Aligners: Sequence Return Aligner (SeqRA) and Stepwise Return Aligner (StepRA).
As previously discussed, DT struggles to align the actual return with the target return due to the under- allocation of attention scores to the return-to-go tokens. One intuitive way to solve this problem is to structure the transformer blocks such that the return-to-go tokens cannot be ignored when processing state and action tokens. To realize this intuition, we introduce Aaluun--Aligeed Decision Transformer (RADT).
We show the model structure of RADT in Fig. We split the input sequence of in Eq. 2 into the return-to-go and other modalities: the return-to to-go sequence т, and the state-action sequence Taa
T, (R1, R.,.... Re).
(7)
Tea (81, a1, 82, .2...., St).
(8)
For practical purposes, RADT processes only the last K timesteps of these sequences. In the transformer block, first apply self-attention to the state-action sequence Tea model dependencies within Tea- We then apply our SeqRA and StepRA to Tsa so that it strongly depends on the return-to-go sequence т,- After the transformer blocks, the action is predicted from the token in the processed state-action sequence Tsa by the prediction head. The model trained using the cross-entropy or mean-squared error loss between the predicted action and the ground truth.
SeqRA and StepRA are designed as complementary alignment strategies that effectively condition the state- action sequence the return-to-go tokens. The first strategy, Sequence Return Aligner (SeqRA), is illus- trated in Fig. 1a SeqRA can capture long-term dependencies involving distant return-to-go tokens by refer- encing multiple past timesteps in the return-to-go sequence. The second strategy, Stepwise Return Aligner (StepRA), is depicted in Fig. 4b StepRA is designed lo capture stepwise dependencies by having each state or action token always reference the corresponding return-to-go token of the same timestep. SeqRA enables the model preserve the influence past returns-to-go, while StepRA ensures an immediate response the current return-to-go. These two strategies can be used independently together. Guided by these strategies, the model naturally acquires a policy that approaches the target return at the end of an episode.

---

Published in Transactions on Machine Learning Research (06/2025)
St-22
Re-1
St
Pr-2
S
Ye-zxPouz
R Re-K+1
R
Re-
(a) Sequence Return Aligner (SeqRA)
(b) Stepwise Return Aligner (StepRA)
Figure 4: Two strategics for conditioning the state-action sequenee on the return-to-go sequence. (a) In Re each stale action token conditioned all return-to-go tokens (Rt-K+1,.... Rÿ), where to capture long-term dependencies. (b) In StepRA, each state or action token is conditioned only the return-to-go token from the same timestep Rj to focus on stepwise relationships.
4.1 Sequence Return Aligner
SeqRA is designed to capture long-term dependencies within the return-to-go sequence Tr, enabling it to identify informative patterns that guide policy adjustments toward long-term objectives. This design lever- ages the sequence modeling capabilities introduced by DT in offiine RL. We employ an attention mechanism to obtain the importance of each token in the return-to-go sequence and integrate the return-to-go sequence into the state-action sequence according to this importance. We make the state-action sequence Tse as the query, and the return-to- sequence T, as both the key and value.
qi = W9,
k2 = TTy Uy = Tr.
W
(9)
These query, key, and value are applied to Eq. 13 to get the attention scores. The attention score ai represents how important the return-to-go token Try is compared to other return-to-go tokens the return-to-go sequence. Note that we use a causal mask ensure that tokens in the state-action se- quence Tea cannot access future return-to-go tokens. By definition, Ej=t-+1 Oij and azj 2 As shown the right side of Fig. T these attention scores are assigned exclusively to return-to-go to- kens, ensuring that the model always references them without allocating any attention to state or tion tokens. According to the allention scores, the return-to-go tokens are aggregated (see Fig. 4a
Oij Trr
(10)
Tye
j=t-K+1
2
SegRA
Linear
The token 2; serves embedding that captures the long-term dependencies of the return-to-go quence in relation to Tsa.i-
Concat
Multi-Head Attention X
1?
We next incorporate the sequence into the state- action sequence Tea- A simple addition may overem- phasize Tea potentially preventing the effective use of the long-term dependencies of the return-to-go sequence from To address this, we learn param- eters that adaptively adjust the scale of against Tsa using powerful method (Nguyen et al.] 2022 from computer vision, which integrates two differ- ent types of features. The flow of this process is illustrated Fig. 5 We concatenate Tea.s and 2; as [& Tso E 2 and obtain dimension-wii wise scaling parameters through learnable affine projection.
State-Action Tokens
Return-to-go Tokens
Figure 5: Illustration of SeqRA. We first apply an attention mechanism to oblain z that integrales the return-to-go sequence T, into the state-action sequence Tea- We then derive a scaling parameter À from Tea and to compute the weighted sum 7

---

Published in Transactions on Machine Learning Research (06/2025)
The addition of 7 and 2, is then formulated as follows, using X,
A W [2: Tsa b.
(11) (12)
20, = (1 + A,)
2,) 8 2; + Tsa.i.
where W RD22 and b E RD are learnable parameters, and denotes the Hadamard product. 7 is the output of SeqRA and the imput to the subsequent StepRA. By zero-initializing W and b, the term 1+\, allows the model to start with basclinc scaling factor of one (simple addition). At the start of training, we can also interpret this as residual connection, a common technique in many transformer based models (Vaswani et al. 2017 Chen et al. 2021 Dosovitskiy et al. 2021 As training progresses, X; is updated, adjusting the balance between 2j and Tsa,i*
4.2 Stepwise Return Aligner
StepRA associates the state action token with the corresponding return-to-go token at the same timestep to capture their stepwise relationship. Specifically, at any timestep j, it applies a linear transformation to the state token Sy or action to- ken aj using weights inferred from the return-to-go token ry. The linear transformation directly em- beds the features of the return-to-go token into the state action tokens. This process is applied sep- arately and identically at each timestep. We train MLP,, and MLP, y to predict the linear transforma- tion weights 7y By E R... The process is formulated as follows:
StepRA
MLP
+
MLP,
Layer Norm
State-Action Tokens
Return-to-go Tokens
Figure 6: Illustration of StepRA. The state and action tokens are first normalized using layer normal-
s) = (1 +L+17 ()) & LayerNorm(SS (3,) + Bg.
(13) ization. They are then linearly transformed using the
af (1 7,) LayerNorm (a,) + Sy
(14) scaling factor y and the shift parameter f, both
(15) inferred from the return token at the same time step.
3/ = (r,, ) 85 = MLP (T;),
where 8y represents the state token Tsa.2j-1, 0j represents the action token Tsa,2j: and ry represents the return- to-go loken Tr.. The process flow of StepRA illustrated in Fig. We first perform layer normalization on the state and action tokens, and then apply this linear transformation, Because layer normalization ensures that the inputs share a uniform scale and distribution, the scaling and shifting parameters learned by the MLPs can work more stably. Similar to 1 + a in Eq. 112. by zero-initializing the parameters of the linear layer, Eqs. (13 and (14] can be considered as the standard layer normalization at the beginning of training. The outputs s and a from the linear transformations are then concatenated to form the sequence r which is used as input to the next components (self-attention, SoqRA, food-forward).
4.3 Arrangement of SeqRA and StepRA
Figure g illustrates the placement of SeqRA and StepRA, which are inspired by the decoder of the original transformer (Vaswani et al.) 2017). SeqRA corresponds to the cross-attention mechanism in the original transformer decoder, while StepRA serves the layer normalization. SeqRA and pross-attention share common function of integrating one type of information with another. In the case of SeqRA, it integrates the return-to-go sequence with the stale-action sequence through an allention mechanism. For this reason, in RADT, we place SeqRA at the position of cross-attention self-attention and food-forward layers). StepRA designed to stabilize training by incorporating layer normalization. Therefore, we position StepRA immediately after each sublayer, including self-attention, SeqRA, and feed-forward.
the original transformer (i.e., between
7

---

Published in Transactions on Machine Learning Research (06/2025)
5 Experiments
We conduct extensive experiments to evaluate our RADT's performance. First, verify that RADT is effective in earning returns consistent with various given larget returns, compared to baselines. Next, we demonstrate through an ablation study that the two types of return aligners constituting RADT are effective individually, and using both types together further improves performance. Furthermore, by comparing the transitions return-to-go throughout an episode, we show that RADT can adaptively adjust its action in response to changes in return-to-go. Finally, we compare RADT with baselines terms of maximizing actual return.
5.1 Datasets
We evaluate RADT on continuous (MuJoCo (Todorov et al. 2012 and discrete (Atari (Bellemare el al. 2013)) control tasks in the same way as DT. MuJoCo requires fine grained continuous control with dense rewards. We use four gym locomotion tasks from the widely-used D4RL (Fu et al. 2020) dataset: ant, hopper, halfcheetah, and walker2d. Atari requires long-term credit assiguments to handle the delay between actions and their resulting rewards and involves high-dimensional visual observations. We use four tasks: Breakout, Pong, Qbert, and Seaquest. Similar to DT, we use 1% of all samples in the DQN-replay datasets per Agarwal et al. 2020) for training.
5.2 Baselines and Settings
We utilize D'T-based methods with various architecture designs as baselines, which enable action generation to be conditioned on the target return. Specifically, we use DT (Chen et al., 2021 p, StARformer (Shang et al. 2022 and Decision ConvFormer (DC) (Kim et al. 2024). For these baselines, we rely on their official PyTorch implementations. Further details about the baselines can be found in Appendix
In MuJoCo, for each method, train three instances with different seeds, and each instance runs 100 episodes for each larget return. In Atari, for each method, 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 trajectorics in the training dataset, specifically from the bottom 5% to the top 5%. This identified range is then equally divided inlo seven intervals, not based on percentiles, but by simply dividing the range into seven equal parts. Fach of these parts represents target return. Further details are provided in Appendix C
5.3 Results
Figures 7 and X show the absolute error between the actual return and target return, plotted for each target return, for MuJoCo and Atari, respectively. We then average the absolute errors over all target returns, and present the mcan and standard error of these averages across seods in Tab. [ for MuJoCo and Tab. 2 for Atari. In these figures and lables, the larget returns, actual returns, and absolute errors are normalized with the largest target return set 100 and the smallest set Overall, RADT significantly reduces the absolute error compared to DT, achieving 44.6% on MuJoCo and 65.5% on Atari, as calculated from Figures 9 and I1 in Appendix A
MuJoCo Domain. In the MuJoCo domain, as shown in Tab. D RADT outperforms all baseline methods across all tasks. Furthermore, as illustrated in Fig. 7 RADT consistently achieves lower absolute error than the baselines in most target returns across all tasks. The baselines exhibit low sensitivity to changes in target returns and tend to bias their actual returns toward specific values. For example, significant errors are observed for baseline methods at lower (arget returns in the ant-mediun-exxprt and walker2d-medium tasks. This is because the ant- mediium-exped and walker2d-mediin datasets contain few trajectories with low target returns. Additionally, Atforrr achieves low absolute error for specific target returns (50 and 66.7) in the ant-medium environment in Fig. 7 However, its overall performance still results in a high absolute error, shown in Tab I In contrast, RADT consistently shows lower absolute errors across wide
8

---

Published in Transactions on Machine Learning Research (06/2025)
Table 1: Absolute errorl comparison in the MuJoCo domain. Target returns are split into seven equally spaced points from the bottom 5% to the top 5% of the dataset. We report the mean and standard error (across three seeds) of the average absolute error over all target returns. comparison for each target return is shown in Fig.
Dataset
Environment
DT
StARformer
DC
RADT (ours)
ant halfcheetah
23.2 = 1.3
14.1
21.5 3.1
3.5 0.5
7.2 4 1.8
23.0 5.0
7.0 0.8
1.8 + 0.6
medium-replay
hopper
12.2 +..2
1.9
0.4
4.4
walker2d
0.1
8.7
1.6
4.0
ant halfcheetah
2.6
2.81
24.9 1.0
9.2 0.3
1..0
2.3
1.6
1.0
medium-expert
hopper
2.2
1..1
8.4
4.8 = 0.5
walker2d
1.72
23.3
17.1
11.0 + 2.3
ant.
37.3 + 2.7
32.8 + 1.2
36.8 3.5
26.3 + 3.0
halfcheetah
36.9 2.8
42.3 0.6
38.8 + 2.0
36.3 - 4.4
medium
hopper
16.3 + 0.3
14.8 + 2.4
13.1 + 1.1
6.5 - 1.1
walker2d
23.6 - 5.7
35.6 2.5
16.5 + 6.4
10.8 = 4.1
halfcheetalh-m-r
hopper-m-r
walker2d-m-r
45
30
18
Absolute
20
2.5
15
10-
6-
25
50
75
100
25
50
75
100
0
25
50
75
100
25
50
75
100
nt=m---e
halfcheetah-m--
hopper-m-e
walker2d-m-e
60
30
40
30
20
Absolute
20
15
10
01
0-
25
50
75
100
25
50
75
100
0
25
50
75
100
25
50
75
100
halfcheetah-m
hopper-m
75/
30
Absolute
50
20
25
25
25
50
75
100
25
50
75
100
25
50
75
100
23
50
75
100
Target Return (Normalized)
DT
StARformer
RADT (ours)
Figure 7: Absolute errort, comparison of across target returns in the MuJoCo domain. We report the mean and standard error over three seeds. The dataset names are shortened: 'mediun-replay" to 'm-r', "medium-ent expertt to 'm-e', and 'medium' to 'm'
range of target returns. These results suggest that RADT can effectively align the actual return with the larget return in environments requiring fine-grained control with dense rewards.
We observe that differences in the environment affect return alignment. From Fig. 10 Appendix A.2 we observe that the actual returns are approximately 75 regardless of the target return across all methods. This is likely because the returns of the trajectories in halfcheetah-medium biased around 75, causing

---

We report the Published in Transactions on Machine Learning Research (06/2025)
Table 2: Absolute errorl comparison in the Atari domain. The way to interpret this table is the same as that of Tab. A comparison for each target return is shown in Fig. g
Game
DT
StARformmer
DC
RADT (ours)
Breakout
21.8 + 5.2
28.8 + 8.9
20.1
11.9 = 3.3
Pong
20.4 + 0.4
10.9 + 3.8
16.6 = 3.9
10.4 + 2.1
Qbert
184.1 + 137.0
176.5 + 103.6
192.7 = 37.2
40.7 = 6.3
Seaquest
42.5 + 18.6
71.0 + 4.4
49.6 + 22.1
21.0 + 4.6
Breakout
Pong
Qbert
Seaquest
80
400
100
Absalute Error
60
40
300
75
40
30
200
50
20
20
100
25
0
25
100
50
75
25
50
100
25
75
100
Target Retur (Normalized)
DT
StARformer
+
DC
RADT (aurs)
Figure 8: Absolute errort comparison across target returns in the Atari domain. mean and standard error over three seeds.
the model overfit to it. Therefore, substantial bias in the dataset's return distribution may reduce the effectiveness of return alignment across most methods.
Atari Domain. In the Atari domain, as shown in Tab. 2 RADT outperforms all baseline methods across tasks. Figure shows that the baselines tend to exhibit large absolute errors for most target returns. This suggests that accurate return modeling the Alari domain is challenging due to the delays bet ween actions and their rewards. Despite this difficulty, RADT consistently demonstrates small absolute errors across wide range of target returns. These results indicate that RADT can effectively match actual returns to target returns in environments requiring long-tern term credit assignment.
However, when examining target returns of 66.7, 83.3, and 100 in Pong, StARformer exhibits the smallest error. We consider this discrepancy to arise from differences in the image encodors used for observations. While RADT uses the same CNN-based image encoder as DT, StARformer employs more powerful Vision Transformer- encoder (Dosovitskiy al. 2021). This suggests that adopting a stronger image encoder could potentially close the performance gap.
5.4 Behavior of RADT in Achieving Return Alignment
In Sec 3 we confirmed that DT cannot adjust its actions according to decreases in the return-to-go, leading to an actual return that significantly exceeds the target return. Tn this section, we compare the behavior of RADT and DT evaluate whether RADT addresses this issuc. For cach target return, we perform 100 runs and calculate the mean and standard error of the return-to-go at each step. Figure plots the transitions of return-to-go input to the models during episode.
For all target returns, DT exhibits nearly lincar decrease in the return-togoo, eventually reaching values lar below zero. This indicates that DT does not adjust its behavior in response to decreases in the return- to-go. These findings suggest that DT generates actions largely independent of the return-to-go decrease. In contrast, RADT demonstrates a gradual docrease in return-to-go that converges near zero, significantly reducing the margin by which exceeds zero compared DT. This suggests that RADT adaptively adjusts its actions in response the changing return-to-go.
10

---

Published in Transactions on Machine Learning Research (06/2025)
Table 3: Absolute error, comparison the ablation study on SeqRA and StepRA. We conduct the ablation study on the medium-replay dataset in MuJoCo. "w/o X" refers the ablation of X from RADT. Each value indicates the mean of absolute errors across multiple target returns as defined in Sec. 5.2 We report the mean and standard error over three sceds, and normalize so that the average and standard error of the DT's results are 1.0 respectively.
Approach
ant
halfcheetah
hopper
walker2d
RADT (ours)
0.15 = 0.42
0.25 + 0.34
0.36 :: 0.13
0.50 + 2.78
w/o SeqRA
0.19 0.43
0.28 0.49
0.59 =
0.56 0.60
w/o StepRA
0.42 2.97
0.32 0.25
0.93 0.73
0.63 3.34
DT
1.00 1.00
1.00 = 1.00
1.00 1.00
1.00 1.00
Table 4: Absolute errort comparison in the ablation study on SeqRA and StepRA in the Atari domain. The way to interpret this table is the same as that of Tab
Approach
Breakout
Pong
Qbert
Seaquest
RADT (ours)
0.55 0.63
0.51 4.72
0.22 + 0.05
0.49 + 0.25
w/o SeqRA
0.61 + 0.60
0.91 + 11.13
0.24 = 0.02
0.95 + 0.72
w/o StepRA
0.57=0.2
0.59 7.16
0.24
0.81 = 1.11
DT
1.00 + 1.00
1.00 1.00
1.00 = 1.00
1.00 + 1.00
5.5 Ablation Study
We conduct ablation studies on the two types of return aligners comprising RADT using the medium- replay dataset in MuJoCo and the Atari dataset. The results for the medium-replay dataset in MuJoCo are summarized in Tab. 8 and those for the Atari dataset are presented in Tab. 4 Each value represents the mean of the absolute errors between the actual return and the target return across multiple target returns. These values are then normalized so that the average and standard error of the DT's results are 1.0. The results from both Tab. 3 and Tab. demonst rate that individually introducing either SeqRA or StepRA is effective. Combining both aligners results in the smallest error across all tasks, implying that SeqRA and StepRA offectively complement each other. In the MuJoCo domain, RADT w/o SoqRA outperforms RADT w/o StepRA, indicating that StepRA is particularly advantageous for continuous control tasks with dense rewards in MuJoCo. Conversely, in the Atari domain, RADT w/o StepRA performs equally or better compared to RADT w/o SeqRA, suggesting that SeqRA is beneficial for long-term credit assignment required in the Atari domain.
5.6 Comparison on Maximizing Expected Returns
In this section, evaluate the influence of our proposed method on maximizing expected returns. We use three environments: MuJoCo tasks characterized by dense rewards, AntMaze tasks featuring sparse rewards, and Atari tasks requiring long-term credit assignment. For RADT, we train three instances with different seeds. Each instance is evaluated over 100 episodes for MuJoCo and AntMaze, and 10 episodes for Atari. The resulting average normalized returns are summarized in Tab. 5 for MuJoCo and AntMaze, and in Tab. 7 for Atari. The normalized returns are computed so that 100 represents the score of an expert policy, as per Fu et al. (2020) for MuJoCo and AntMaze, and Hafner et al. (2021) for Atari. Additionally, these tables include the highest episode returns in each dataset.
From Tab. 5 and Tab. G RADT achieves performance close to the maximum returns of the datascts in MuJoCo and AntMaaee, and its performance is comparable lo or exceeds that of DT and DC. Similarly, from Tab. 7 RADT reaches returns near the maximum values of the datasets in Breakout and Pong. Qbert and Seaquest, although RADT also successfully achieves returns well above the maximum returns of the datasets, it falls short of the baselines specializing in maximizing returns, despite outperforming all approaches in return alignment shown in Tab 2 RADT effectively reproduces near-maximum dataset
11

---

Published in Transactions on Machine Learning Research (06/2025)
Table 5: Performancef comparison of return maximization in the MuJoCo domain. We report the average across three seeds from our simulation results for RADT. The boldface numbers denote the maximum score. We cite the results for DT and DC from their reported scores. We exclude StARformer from the comparison since the original paper does not report results MuJoCo tasks. The max return of the dataset indicates the highest episode returns for each dataset.
Max Return of Dataset
Datasel
Environment
DC
DT
RADT (ours)
halfcheetah
41.3
36.6
41.3 0.30
42.4
modium-roplay
hopper
94.2
82.7
95.7 0.22
98.7
walker2d
76.6
66.6
75.9 + 1.55
90.0
halfcheetah
93.0
86.8
93.1 + 0.01
92.9
medium-cxpert
hopper
110.4
107.6
110.4 + 0.38
116.1
walker2d
109.6
108.1
109.7 0.16
109.1
Table 6: Performance" comparison of return maximization the AntMaze domain. We cite the results for DC from their reported scores. We exclude StARformer and DT from the comparison since the original papers do not report results AutMaze tasks. The way to interpret this table the same that of Tab.
Max Return of Dataset
Dataset
DC
RADT (ours)
umaze
85.0
90.7 7 + 4.35
100.0
umaze-diverse
78.5
80.7 + 2.37
100.0
Table 7: Performancet comparison of return maximization in the Atari domain. We cite the results for DC, StARformer, and DT from their reported scores. The way interpret this table the same that of Tab.
Max Return of Dataset
Environment DC
StARformer
DT
RADT (ours)
Breakout
352.7
436.1
267.5
302.6 + 22.6
371.4
Pong
106.5
110.8
106.1
107.0 + 2.1
116.7
Qbert
67.0
51.2
15.4
12.8 3.5
5,8
Seaquest
2.7
1.7
2.5
0.1 0.6
performance, but faces challenges in scenarios requiring extrapolation to substantially higher target returns beyond the training distribution, We further discuss this limitation in Sec.
6
Related work
Return-condditoondd Offline RL Recent studies have focused on formulating offline reinforcement learn-
ing (RL) problem of predicting action sequences that are conditioned by goals and rewards (Chen al.
2021
Janner et al. 2021 Emmons al. 2022 David et al. 2023 SSchiidhuber
2019
Srivastava et al.
2019
Fujimoto & Gu Kostrikov et al. 2022) by modeling the relationship between rewards and actions through supervised
This approach differs from the popular valuc-based methods (umarr et al. 2020
2021
learning. Decision transformer (DT) (Chen et al. 2021 introduces the concept desired future returns and improves performance by training the transformer architecture (Vaswani al. 2017) as a return-conditioned
policy. Based on DT, various advancements have been proposed for introducing value functions (Yamagata)
et al.
2023
Gao et. al. 2024 Wang al.] 2024), finetuning models online (Zheng et. al. 2022 adjusting
12

---

Published in Transactions on Machine Learning Research (06/2025)
the history length (Wu al. 2023b), and improving the transformer architecture (Shang et al. 2022 Kim et al. 2024). However, these approaches do not explicitly address the alignment between the actual returns and the target return, which is a central focus of our work.
Improving Transformer Architecture for Offline RL Among the advancements based on DT, we delve deeper into studies that focus on modifying the model architecture similarly to our approach. Some efforts focus on refining the transformer architecture for offine RL. StARformer (Shang et al. 2022) in- troduces 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 leeding 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 convoltion to capture the inherent. local dependence pattern of MDP. These methods focus primarily on improving aspects of performance. To the best of our knowledge, this is the first work to explicitly align actual returns with target returns as the main objective in return-conditioned offiine RL.
7 Discussion and Conclusion
In this paper, we proposed RADT, a novel decision-making model for aligning the actual return with the target return in offline RL. RADT splits the input sequence into return-to-go and state-action sequences, and reflects return-to-go in action generation by uniquely handling the return-to-go sequence. This unique handling includes two stralegies that capture long-term dependencies and stepwise relationships within the return-to-go sequence. Experimental results demonstrated that RADT has superior aligning capabilities compared to existing DT-based models. One limitation of our method is a slight increase in computational cost compared to DT. This could potentially be improved by introducing a lightweight attention mechanism, such as Flash Attention (Dao al. p022), ; into our S.qRA. Another limitation is the difficulty adapting target returns that significantly exceed the maximum returns within the dataset. This issue could potentially be addressed by enhancing the training data returns through data augmentation techniques Li et al. (2024); Luo et al. (2025). We believe that RADT's alignment capability will improve the usability of offine RL agents in a wide range of applications: creating diverse AI opponents in video games and educational tools, controlling heterogeneous agents in simulations, and efficiently generating human motions for animation or robotics. Investigating how RADT generalizes to multi-agent environments with diverse and adaptive opponents is an exciting avenue for future work.
Acknowledgments
This work was supported by JST, ACT-X Grant Number JPMJAX23CE, Japan. Kaito Ariu is supported by JSPS KAKENHI Grant No. 23K19986
References
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on oflline rein- forcement learning. In Proc. of 1CML, 2020.
Mare G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47(1):253 279, 2013.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Tee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Proc. of NewrIPS, 2021.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-cfficient exact allention with io-awareness. In Proc. of NewrIPS, 2022.
Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision S4: Efficient sequence- based RL via state spaces layers. In Proc. of ICLR, 2023.
13

---

Published in Transactions on Machine Learning Research (06/2025)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale. In Proc. of ICLR, 2021.
Scott Emmons, Benjamin Eysenbach, Nya Kostrikov, and Sergey Levine. Rys: What essential for offline RL via supervised learning? In Proc. of ICLR, 2022.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data- driven reinforcement learning, 2020.
Scott Fujimoto and Shixiang Gu. A minimalist approach to offline reinforcement learning. In Proc. of NewrIPS, 2021.
Chenxiao Gao, Chenyang Wu, Mingjun Cao, Rui Kong, Zongzhang Zhang, and Yang Yu. Act: Empowering decision transformer with dynamic programming via advantage conditioning, In Proc. of AAAI, 2024.
Danijar Hafner, Timothy Lillierap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In Proc. of ICLR, 2021.
Kaidong Hu, Brandon Haworth, Glen Berseth, Vladimir Pavlovic, Petros Faloutsos, and Mubbasir Kapa- dia. Heterogeneous crowd simulation using parametric reinforcement learning. IEEE Transactions on Visualization and Computer Graphics, 2023.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Proc. of NewrIPS, 2021.
Zhengyao Jiang, Yingchen Xu, Nolan Wagener, Yicheng Luo, Michael Janner, Edward Grefenstette, Tim Rocktäschel, and Yuandong Tian. H-GAP: Humanoid control with generalist planner. In Proc. of ICLR, 2024.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline RL? In Proc. of ICML, 2021.
Jeonghye Kim, Suyoung Lee, Woojun Kim, and Youngchul Sung. Decision ConvFormer: Local filtering MetaFormer sufficient for decision making. In Proc. of ICLR, 2024.
lyya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning, In Proc. of ICLR, 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline rein- forcement learning. In Proc. of NewrlPS, 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems, 2020.
Guanghe Li, Yixiang Shan, Zhengbang Zhu, Ting Long, and Weinan Zhang, Diffstitch: Boosting offline reinforcement learning with diffusion-based trajectory stitching. In Proc. of ICML, 2024.
Yunhao Luo, Utkarsh A Mishra, Yilun Du, and Danfei Xu. Generative trajectory stitching through diffusion composition, 2025.
C.M. Macal and M.J. North. Tutorial on agent-based modeling and simulation. In Proc. of WSC, 2005.
Van-Quang Nguyen, Masanori Suganuna, and Takayuki Okatani. Grit: Faster and better image captioning transformer using dual visual features. In Proc. of ECCV, 2022.
Andreas Panayiotou, Theodoros Kyriakou, Marilena Lemonari, Yiorgos Chrysanthou, and Panayiotis Char- alambous. Cop: Configurable crowd profiles. In Proc. of SIGGRAPH, 2022.
14

---

Published in Transactions on Machine Learning Research (06/2025)
Xue Bin Peng, Yunrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. Ase: Large-scale reusable adversarial skill embeddings for physically simulated characters. ACM Trans. Gruph., 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, Lya Sutskever, et al. Improving language understanding by generative pre-training, 2018.
Davis Rempe, Zhengyi Tuo, Xue Bin Peng, Ye Yuan, Kris Kitani, Karsten Kreis, Sanja Fidler, and Or Litany. Trace and pace: Controllable pedestrian animation via guided trajectory diffusion. In Proc. of CVPR, 2023.
Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards just map them to actions, 2019.
Jinghuan Shang, Kumara Kahatapitiya, Xiang Li, and Michael S. Ryoo. StARformer: Transformer with stale-action-reward representations for visual reinforcement learning. In Proc. of ECCV, 2022.
Ruimin Shen, Yan Zheng, Jianye Hao, Zhaopeng Meng, Yingleng Chen, Changjie Fan, and Yang Liu. Generating behavior-diverse game ais with evolutionary multi-objective deep reinforcement learning. In Proc. of IJCAI, 2021.
Adish Singla, Anna N Rafferty, Goran Radanovic, and Neil T Heffernan. Reinforcement learning for educa- tion: Opportunities and challenges, 2021.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning, 2019.
Emanuel Todorov, Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In Proc. of IROS, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. of NeurIPS, 2017.
Yuanfu Wang, Chao Yang, Ying Wen, Yu Liu, and Yu Qiao. Critic-guided decision transformer for offline reinforcement learning. In Proc. of AAAT, 2024.
Shuang Wu, Jian Yao, Haobo Fu, Ye Tian, Chao Qian, Yaodong Yang, QIANG FU, and Yang Wei. Quality- similar diversity via population based reinforcement learning. In Proc. of ICLR, 2023a.
Yuch-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. In Proc. of NeurIPS, 2023b.
Yiming Xic, Varun Jampani, Lei Zhong, Deqing Sun, and Huaizu Jiang. Omnicontrol: Control any joint at any time for human motion generation. In Proc. of ICLR, 2024.
Haoran Xu, Li Jiang, Jianxiong Li, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. In Proc. of NeurIPS, 2022.
Taku Yamagala, Ahmed Khalil, and Raul Santos-Rodright Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline RT. In Proc. of TCML, 2023.
Georgios N. Yannakakis and Julian Togelius. Artificial Intelligence and Games. Springer, 2018
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. In Proc. of NeurIPS, 2021.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In Proc. of ICML, 2022.
W. Zhu, X. Ma, D. Ro, H. Ci, J. Zhang, J. Shi, F. Gao, Q. Tian, and Y. Wang. Human motion generation:
A survey. IEEE Transactions on Pottern Analysis and Machine Intelligence, 2024.
15

---

Published in Transactions on Machine Learning Research (06/2025)
Broader Impact
RADT can bring about positive social impact by enabling adaptation to new application fields such gaine production, educational content production, and simulations, as it can control the performance of agents more accurately. However, this advantage also comes with potential negative impacts such as the tracing of user behavior patterns. Such impact can be mitigated by applying methods such as blinding personal information during data generation and collection process.
A Additional Experimental Results
To better understand RADT, this section first explains additional analysis of the discrepancies between actual return and target return. Then, describe an ablation study on adaptive scaling from Sec 4.1 and ablation studies on each sublayer in the Atari domain.
A.1 Additional Analysis of Discrepancies
21.9
19.5
Absolute Error
17.7
10.8
DT
DC
StARformer
RADT (ours)
Figure 9: Absolute errors, between the target return and the actual return in the MuJoCo domain.
We present Fig. 9 which summarizes the discrepancies between the actual return and the target return in the MuJoCo domain. The figure shows the averages of the absolute errors between the actual return and the target return from Fig. 7 across seven target returns and twelve tasks. We normalize these absolute errors by the difference between the top 5% and bottom 5% returns in the dataset 3 RADT outperforms the baselines, achieving a 44.6% reduction in discrepancies compared to DT. Among the baselines, DC has the smallest discrepancies, while the StARformer has the largest.
To further delve into the results of the MuJoCo domain, we present comparison of the actual returns in Fig 10 The black dotted line represents y æ, indicating that the actual return matches the target return perfectly. The closer to the black dotted line, the better the result. In all tasks except halfchectah- medium, RADT is closer to the target return than the baseline is. It can be seen that ant-medium and halfcheetab-mmedium medium are struggling due to the extremely biased distribution of target returns in the datasets. In some tasks, the baselines show constant actual return regardless of the input target return (e.g., DT in modtmm-uuu---ppay, StARforrmr in walkor2d, DC and StARformer in wem-immmdx-perttttt etc.). We believe this is due lo the models overfitting the target return in where the data is concentrated.
We present Fig I1 which illustrates the discrepancies between the actual return and the target return in the Atari domain. The values in this figure represent the average absolute errors between the actual return and
We normalize returns from each datasct split, unifying the normalized ranges facilitate comparisons between splits (medium, medium-replay. medeum-experrt). DaRL's normalization compare results because the achievable performance limits trained agents vary significantly between splits.
same values across different splits, making difficult.
16

---

Published in Transactions on Machine Learning Research (06/2025)
ant-m-r
halfcheetah-m--
hopper-m-r
walker2d-m-4
100
100
120
100
75
75
50
60
50
25
30
25
0
25
50
75
100
0
25
50
75
100
0
75
100
25
50
75
100
ant-m-e
halfcheetah-m--e
hopper-m-e
walker2d-m-e
100
100
100
100
75
75
75
75
50
50
25
25
0
25
50
100
100
ant-m
halfcheetah-m
hopper-m
walker2d-m
100
100
100
100
75
50
25
25
25
0
0
25
50
75
100
0
25
100
0
25
75
100
25
50
75
100
Target Return (Noraaiized)
DT
+
StARformer
DC
RADT (ours)
Figure 10: Comparisons of actual returns per target return in the experiments of Fig. 7 for MuJoCo. Each column represents task. The x-axis represents the target return, and the y-axis represents the actual return. The x-axis and y-axis are normalized in the same way as in Fig. 7 Targel returns are set in the same way as 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. The dataset names are shortened:'ment "mediumrrpplay" 'm-r', ceddumm-expert 'm-e', and 'medium' 'm'.
17

---

Published in Transactions on Machine Learning Research (06/2025)
79.2
74.6
74.3
Absolute Error
25.7
DT
DC
StARformer
RADT (ours)
Figure 11: Absolute crrors. between the target return and the actual return in the Atari domain.
the target return, averaged over seven target returns and four tasks. These absolute errors are normalized by the difference between the top 5% and the bottom 5% of returns in the dataset. RADT significantly outperforms the baseline methods. Specifically, it achieves 65.5% reduction in discrepancies compared to DT, which is greater reduction margin than what we observed in the MuJoCo domain. Among the bascline methods, DC exhibits the smallest discrepancics, while StARformer shows the largest, pattern that is consistent with the MuJoCo domain.
A.2 Additional Ablation Study
We conduct an ablation study on the adaptive scaling in SeqRA on the medium-replay dataset in MuJoCo and the Atari datasct. The experimental sctup is the same as described in See 5.5 The results are summa- rized in Tab.B and Tab. g These results demonstrate that introducing adaptive scaling (RADT w/o StepRA) outperforms the case where only the attention mechanism of SeqRA is used (RADT w/o StepRA and Adap- tive Scaling). This suggests that adaptive scaling effectively merges which includes information from the return sequence, with the state-action sequence Tea-
Table 8: Absolute errorl, comparison in the ablation study for the adaptive scaling in SeqRA on the medium-replay dataset in MuJoCo domain. The way to interpret this table is the same that of Tab
Approach
ant
halfcheetah
hopper
walker2d
RADT (ours)
0.15 + 0.42
0.25 + 0.34
0.36 + 0.24
0.50 + 2.78
w/o StepRA
0.42 + 2.97
0.32 + 0.25
0.93 + 0.73
0.63 + 3.34
w/o StepRA and Adaptive Scaling
0.74 + 2.94
0.85 + 0.45
0.99 = 0.73
0.87 + 2.00
DT
1.00 1.00
1.00 + 1.00
1.00 + 1.00
1.00 + 1.00
Table 9: Absolute errorl comparison in the ablation study for the adaptive scaling in SeqRA on the Atari domain. The way to interpret this table is the same as that of Tab
Approach
Breakout
Pong
Qbert
Seaquest
RADT (ours)
0.55 + 0.63
0.51 4.72
0.22 0.05
0.49 0.25
w/o StepRA
0.57 + 0.20
0.59 + 7.16
0.04
0.81 + 1.11
w/o StepRA and Adaptive Scaling
0.72 + 1.52
0.86 2.16
0.26 0.07
0.86 + 1.22
DT
1.00 + 1.00
1.00 + 1.00
1.00 + 1.00
1.00 + 1.00
18

---

Published in Transactions on Machine Learning Research (06/2025)
A.3 Additional Analysis of Parameter Size
We compare RADT and models where we increase the number of blocks in DT to bring the number of parameters closer to RADT. Table T0 shows the parameters for each model on the medium-replay dataset of MuJoCo. Because RADT adds both SeqRA and StepRA to the transformer block, it has 1.6 times the number of parameters compared to DT. By increasing the number of blocks in DT to six, it surpasses the number of parameters of RADT. Therefore, we compare RADT with a DT that has six blocks and present the results on the medium-replay dataset of MuJoCo in Tab. RADT outperforms the DT with six blocks. On the other hand, we find that DT with 6 blocks performs worse than that with blocks the hopper and walker2d environments. This result suggests that simply increasing the number of parameters of DT does not improve return alignment, and that the improved model structure of RADT is effective.
Table 10: Comparison of parameters. Since the dimensions of the state and action spaces vary for each dataset, the model size differs depending on the dataset.
Environment D3 blocks)
DT(5 blocks)
DT blocks)
RADT(3 blocks)
ant halfcheetah
740,232
1,136,776
1,335,048
1,219,848
727,686
1,124,230
1,322,502
1,207,302
hopper 146,147
1,122,691
1,320,963
1,205,763
walker2d
727,686
1,124,230
1,322,502
1,207,302
Table 11: Absolute errorl comparison between RADT and DT with increased parameters. We use the medium-replay datasot MuJoCo. The way to interpret this table is the same as that of Tab
Environment
DT(3 blocks)
DT blocks)
RADT(3 blocks)
ant
23.2 + 1.3
16.8 = 4.5
3.5 + 0.5
halfcheetah
7.2 + 1.8
5.5 + 0.5
1.8 + 0.6
hopper
12.2 + 4.3
12.8 + 2.4
4.4 0.5
walker2d
8.1 + 0.3
10.0 + 0.7
4.0 + 0.8
B
Baseline Details
We use the model code for DT, StARformer, and DC from the following sources. DT: https://github. comm/eel/@ecisiionttaaantttrer StARformer: https://github. httpaaa....oeellaaaaaaaaaatt....
DC: Although StARformer uses step-by-step rewards instead of returns, in our experiments, we employ return-conditioning using returns. This modification allows St.A- former to condition action generation on target return. The original paper (Shang et al. 2022) states that this modification has a minimal impact on performance. 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-basod embeddings like Vision Transformer (Dosovitskiy et al. 2021).
The baseline results for aligning the actual return with the target return (Sec. 5.3 and the ablation study (Sec.55. 5.5) are from our simulations. The hyperparameters for each method in our simulations are set according to the defaults specified in their original papers or open-source codebases. The bascline results for maximizing the expected return (Sec. 5.6 stem from the original papers or third-party reproductions.
C Experimental Details
For the reproducibility of the experiments, this section explains the comparison of computational costs and the hyperparameters used in the experiments.
19

---

Published in Transactions on Machine Learning Research (06/2025)
Table 12: Comparision of computational cost.
Method
Training Time (s)
GPU memory usage (GilB
DT 363
0.030
RADT
466
0.034
C.1 Comparison of Computational Cost
Table 12 shows comparison of the computational costs of DT and RADT. We compare the training time and GPU memory usage incurred when running 10 iterations of training on the hopper-medium-replay dataset. In this comparison, we use an NVIDIA A100 GPU. RADT has slight increases in computation time and memory usage from DT. We believe these increases are due to the addition of SeqRA and StepRA. The computational costs RADT could potentially be improved by introducing efficient attention mechanisms such as Flash Attention [ao al. 2022
C.2 Hyperparameters
Our implementation of RADT based on the public codebase of DT. We used an Nvidia A100 GPU for training in the Atari and MuJoCo domains. The full list of hyperparameters of RADT is found in Tab. 13
and Tab. 14 The hyperparameter settings of RADT are the same in both aligning and maximixing.
Table 13: Hyperparameters settings of RADT the MuJoCo domain and the AntMaze domain. The dataset names are shortened: 'medium" 'm', 'medium-expert! to 'm-e', 'umaaze' to 'u', and "umaze- diverse' to 'u-d.
Hyperparameter
Value
Number of blocks
2, hopper-m
3, otherwise
Number of heads
256, ant-m-e, halfcheetah-m,
Embedding dimension
antmaze-u, antmaze-u-d
128, otherwise
Batch size
256, antmaze-u, antmaze-u-d
64, otherwise
Nonlinearity function
GELU, transformer
SILU, StepRA
Context length
20
Dropout
0.1
Learning rate
10-4
Grad norm clip
0.25
Weight decay
10-4
Learning rate decay
Linear warmup for first 10ah steps
D Additional Real-World Application Example
We describe the additional real-world application that requires return alignment, mentioned in Sec.
Example (Agent-based modeling and simulation). Agent-based modcling and simulation (Macal North 2005) involves examining system dynamics by representing them collections of interacting, autonomous agents. In this context, incorporating heterogeneous behaviors (Hu et al. 2023 and controllable poli-
20

---

Published in Transactions on Machine Learning Research (06/2025)
Table 14: Hyperparameters settings of RADT in the Atari domain.
Hyperparameter
Value
Number of blocks
Number of heads
Embedding dimension
128
Batch size
512 Pong
128 Breakout, Qbert, Seaquest
Context length K
50 Pong
Breakout, Qbert, Seaquest
Nonlinearity
ReLU encoder
GELU transformer
SiLU StepRA
Encoder channels
32, 64, 64
Encoder filter size
8 8,4
Encoder strides
4,2,1
Max epochs
Dropout
0.1
Learning rate
6 10
Adam betas
(0.9, 0.95)
Grad norm clip
1.0
Weight decay
0.1
Linear warmup and
Learning rate decay
cosine decay
Warmup tokens
512 20
Final tokens
500000 K
cics Pannayiotou et al. 2022 is essential. Our approach introduces controlled variety of agent behaviors by adjusting their performance levels based on a specified larget return.
21


Reviews:
[{"review_id": "E8iyHkSFUw", "reviewer": "Reviewer_914B", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is clearly motivated (especially Fig 3 showcases the addressed problem well)\n- The aligner modules are intuitive \n- The evaluation is solid since it compares multiple state-of-the-art approaches on common benchmarks and shows the improvement of achieving specified target returns of the RADT\n\nWeaknesses:\n- The motivation and Sec. 3 could be shortened\n- Limitations and the differences between different benchmarks should be discussed in more detail", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposed an adapted transformer architecture for decision transformers such that improves the capability of exactly achieving target returns instead of higher returns. For that, two aligners are proposed. The new architecture is evaluated on Atari and Mujoco benchmarks for multiple trajectory data sets. The evaluation shows a consistent performance improvement compared to other state-of-the-art approaches on the selected benchmarks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Example 2 and 3 in Sec. 1 are very similar and I would suggest removing one of them \n- Fig. 2 was not easy to interpret for me at that point of the paper and it is unclear from which experiments the values stem. I would suggest making the figure more stylistic, making the caption clearer on where these numbers stem from, or removing the figure\n- Sec. 3 could be significantly shortened, so that it briefly states the hypothesis and explains the figure.\n- I would have expected that low target returns are easier to achieve than high, but it looks like that if there is a difference within a benchmark for the RADT performance, it is mainly for lower returns. Do you have an intuition for that? \n- For the game example from the introduction, the performance of an RADT agent playing a human player depends on their interaction. Two beginner players could be lacking in two very different ways, so the RADT agent is only a good competitor for one of them. The benchmarks you consider are single-agent benchmarks. Do you think your approach generalizes to multi-agent setups well? What could be good first benchmarks to investigate such setups?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0qgo1yAc5I", "reviewer": "Reviewer_2ThK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n\n1. When I first started reading this paper, I found it difficult to think of a scenario where a suboptimal policy would be needed. However, the three examples addressed all my concerns. The authors also clearly explain the limitations of DT and motivate their idea effectively.\n\n2. The paper’s main claim is clear and well-supported by the experimental results.\n\n**Weaknesses**\n\nNo clear weaknesses.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new attention model for offline RL, treating offline RL as a sequential modelling problem. Compared to prior work, the model gives higher attention to the return-to-go, effectively aligning the actual return with the target return-to-go. Additionally, it achieves better performance in many MuJoCo environments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. In the second paragraph of Section 3, the paper states that “when the majority of the attention scores are allocated to state or action tokens, the allocation to return-to-go tokens decreases.” It would be helpful if the authors could further explain why the attention scores tend to concentrate on state or action tokens. Why does the return-to-go token receive inadequate attention in this setting?\n\n2. I would suggest that the authors improve the completeness of the experiments. Since the paper’s claim is direct and does not require overly complex analysis, it would be helpful to compare the performance results (Section 5.6) and conduct the ablation study (Section 5.5) in both the MuJoCo and Atari domains, using a wider range of baselines.\n\n3. The authors should improve Section 6 by better reflecting current research methods and emphasising how their approach differs. Are there other methods that modify the model architecture to improve aspects of performance in modelling sequence RL? Is this the first paper to explicitly align predicted returns with return-to-go? Rather than simply listing related work, it would be more informative to provide a critical comparison and highlight the unique contributions of this paper.\n\n4. I don’t think Figure 1 is necessary. The text description is clear enough.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PfltJiNuQJ", "reviewer": "Reviewer_TojF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. The paper proposes a simple and intuitive architectural change to force the model to incorporate reward information more directly. \n\n2. Experiments verify that indeed the DT model is often ignoring the conditioned returns and thus overshooting the desired return.\n\n3. Experiments clearly show large improvements in terms of the calibration of RADT in the sense that the actual returns are much closer to the desired conditioned returns across environments.\n\n4. Ablation experiments illustrate that both sequence-level and step-level conditioning are useful components of RADT.\n\n5. The presentation is generally clear.\n\n## Weaknesses\n\n1. One issue is that the proposed architecture changes add parameters within every block. As a result the RADT architecture is larger than the DT baselines. This is mentioned briefly in the appendix discussing computational time and memory requirements. However, it would seem to be important to add a parameter-matched baseline DT with slightly larger model so that it matches the parameters in the RADT model.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a modification to the architecture of a decisiton transformer (DT) so as to put more focus on the rewards/returns at each layer called RADT. By doing so, the model improves in terms of the error of the desired return compared to the actual return. Experiments verify this on mujoco and atari domains.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Overall, I think this is a very solid paper with a clear contribution and relatively thorough experimentation. I would just like to see the one extra experiment listed above for completeness.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "02oSfj33vD", "reviewer": "Reviewer_i1BT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n* The problem is well-motivated and related to the subject of controllable agents.\n* The finding of the attention mechanism in the vanilla decision transformer being a possible main contributor to misaligned returns is a good one.\n\n**Weakness/questions**\n* The motivations for the proposed architecture modifications were not super clear. It seemed that the only goal was to ensure reward-related information is routed to the output head.\n* Intuitively in RL via SL, data quality/diversity is important. We would expect models trained on data with diverse returns to be better return-aligned. Why do we see medium-replay having better overall alignment? Is that because it has more diverse returns?\n* Why in certain environments the return error can be as high as 75 (e.g., halfcheetah-m)? I think it'd be good to comment on the effect of the environments on the results, e.g. the difficulty of achieving a specific level of return is not necessarily in every environment.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a few modifications to the decision transformer architecture in the context of return-conditioned reinforcement learning such that rewards achieved by the agent aligns more closely with the rewards it is conditioned on. The main modifications are \"sequence return aligner\" which introduces additive interactions with state-action tokens and \"stepwise return aligner\" which introduces multiplicative interactions with state-action tokens. The experiments show that the proposed approach significantly improves return alignment accuracy in Mujoco and Atari environments compared to baselines (including the vanilla decision transformer).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I don't have other requested changes.", "broader_impact_concerns": "There's no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers recommended acceptance. There is a consensus that this paper identifies a meaningful research question, proposes a well-designed solution, and presents sufficient empirical evidence to demonstrate its effectiveness. It is a solid empirical contribution, aligning well with TMLR’s emphasis on technical soundness over subjective significance or novelty. As stated earlier, some minor points deserve improvement, hence the decision to *accept with minor revision*.
Decision: Accept (poster)