Peak-Return Greedy Slicing: Subtrajectory Selection for Transformer-based Offline RL
Authors: Zhiwei Xu
Miduo Cui
Dapeng Li
Zhihao Liu
Haifeng Zhang
Hangyu Mao
Guoliang Fan
Bin Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Offline Reinforcement Learning
Transformer
Primary area: reinforcement learning
TL;DR: We propose Peak-Return Greedy Slicing (PRGS), a simple yet effective framework that enhances the stitching ability of Transformer-based offline RL by explicitly selecting high-return subtrajectories at the timestep level.

Abstract: Offline reinforcement learning enables policy learning solely from fixed datasets, without costly or risky environment interactions, making it highly valuable for real-world applications. While Transformer-based approaches have recently demonstrated strong sequence modeling capabilities, they typically learn from complete trajectories conditioned on final returns. To mitigate this limitation, we propose the Peak-Return Greedy Slicing (PRGS) framework, which explicitly partitions trajectories at the timestep level and emphasizes high-quality subtrajectories. PRGS first leverages an MMD-based return estimator to characterize the distribution of future returns for state-action pairs, yielding optimistic return estimates. It then performs greedy slicing to extract high-quality subtrajectories for training. During evaluation, an adaptive history truncation mechanism is introduced to align the inference process with the training procedure. Extensive experiments across multiple benchmark datasets indicate that PRGS significantly improves the performance of Transformer-based offline reinforcement learning methods by effectively enhancing their ability to exploit and recombine valuable subtrajectories.

Published as a conference paper at ICLR 2026
PEAK-RETURE GREEDY SLICING: SUBTRAJECTORY SELECTION FOR TRANSFORMER-BAST -BSED OFFLINE RL
Zhiwei Xu' Miduo Cui', Dapeng Li Zhihao Liu Haifeng Zhang², Hangyu Maoe,
Guoliang Fan", Bin Zhang?*
School of Artificial Intelligence, Shandong University, Jinan, China
2 The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, China
Li Auto Inc.. China
4 T Tsinghua University, Beijing, China
5 Zhongguancun Academy, Beijing, China
Institute of Microelectronics, Chinese Academy of Sciences, Beijing, China zhiwei.xuu xu@ sdu. edu. cn, zhragbin20200ia.ca
ABSTRACT
Offline reinforcement learning enables policy learning solely from fixed datasets. without costly or risky environment interactions, making highly valuable for real-world applications. While Transformer-basee approaches have recently demonstrated strong sequence modeling capabilities, they typically learn from complete trajectories conditioned on final returns. To mitigate this limitation, we propose the Peak-Return Greedy Slicing (PRGS) framework, which explic- itly partitions trajectories at the timestep level and emphasizes high-quality sub- trajectories. PRGS first leverages an MMD-base return estimator to characterize the distribution of future returns for state-action pairs, yielding optimistic return estimates. then performs greedy slicing to extract high-quality subtrajectories for training. During evaluation, an adaptive history truncation mechanism is in- troduced to align the inference process with the training procedure. Extensive experiments across multiple benchmark datasets indicate that PRGS significantly improves the performance of Transformer based offline reinforcement learning methods by effectively enhancing their ability to exploit and recombine valuable subtrajectories.
INTRODUCTION
In practice, offline reinforcement learning (RL) (Levine al. 2020), which learss policies purely from fixed datasets without environment interaction, has garnered significant attention in recent years. Unlike online RL, offline RI. eliminates the need for costly or potentially risky interactions with the environment, making it particularly valuable in domains such as autonomous driving (Lin et al. 2024), robotics (Kumar al.) 2023). and recommender systems (Xin et al., 2022). With advances in sequence modeling, Transformer-baue based offline RL (Chen et al. 1202 Janner et al.) 2021 has gained increasing traction. Leveraging the Transformer's strength modeling long-range dependencies and its expressive power (Vaswani et al. 2017), recent methods have achieved sub- stantial improvements on several benchmark datasets, presenting novel paradigm for the progress of offline reinforcement learning.
However, existing Transformer-ban based offline RI. methods remain limited in effectively stitching to- gether high-quality segments from different trajectories (Branndfonbrener et al. 2022). The stitching capability refers to the algorithm's ability to identify and recombine superior trajectory fragments, thereby enabling policy learning that outperforms what can be achieved from any single existing trajectory. Although some prior works have explored trajectory resampling (Li et al. 2024 Lec et al. 2024), value- based guidance (Yamagata et al.) 2023
Pei et al.
2025
Wang al. (024).
conditional modeling (Zhang et al.) 2024 Kim et al.
2024
Wu et al.
2023 to mitigate this issue,
"Corresponding author.

---

Published as a conference paper at ICLR 2026
their effectiveness is constrained by the coarse granularity of trajectory-level processing. As result, these methods often struggle to handle data with uneven quality and fall short of fully leveraging the power of the Transformer.
is worth noting that humans, when learning to make decisions, do not evaluate experience solely based on its final outcome. Instead, they tend to distinguish good subtrajectories and bad sub- trajectories within long trajectories. As illustrated in Figure 1 even if a trajectory leads to a suboptimal outcome overall, it may still contain locally high-value experiences. Current Transformer- based methods generally derive insights from complete trajectories, fo- cusing exclusively on their final returns. In contrast, humans se- lectively retain valuable subtrajectories and enhance learning effi- ciency by composing new experiences through stitching (Schank Abelson 2013; Nakahashi et al., 2016 Dennett 1989; |Schachner & Carey 2013). This capability of subtrajectory selection is notably absent Transformer-basee offline RL methods, primarily because identifying appropriate split points within a trajectory is often challenging.
Bad
Subtrijectories
subtrajectories
Split Point
Figure 1: Illustration of tra- jectory slicing with an appro- priate split point.
Inspired by this observation, propose the hance the performance of Transformer- based offline RI. methods by enabling them to identify and focus on high-quality experience subtrajectories. Concretely, we first train return estimator based on Maximum Mean Discrepancy (MMD) (Smola et al.) 2006). which approximates the distribution of potential future returns for each state-action pair under optimistic assumptions. Leveraging this estimator, PRGS infers the initial-to-go return for each state along a trajectory and recursively seg- ments the entire trajectory into multiple subtrajectoriees Positions with peak return are identified as splir points. This fine tne-grained slicing at the timestep level allows PRGS to effectively preserve valuable experience segments and stitch them. Finally, to ensure consistency in history length be- tween training and evaluation, we introduce an adaptive mechanism that truncates the input history dynamically during evaluation. To evaluate the effectiveness of PRGS, we integrate it into several Transformer-baseee offline RL algorithms and conduct comprehensive experiments across a variety of benchmark domains. Experimental results demonstrate that PRGS may be seamlessly integrated into existing frameworks and consistently yields average performance improvement of 15.8% over the original baselines across diverse tasks.
Greedy Slicing (PRGS) framework to
2
PRELIMINARY
2.1 OFFLINE REINFORCEMENT LEARNING
In offline RL, agent is trained solely on static dataset D = [T]. Each trajectory r == {(s, a, r))
consists of state s E S, action a E A. and reward r E R at each time step. The objective is to learn policy π.: S AA that maximizes the expected discounted return J(m) Er-= EE 2'rz]. where E (0, 1] is the discount factor. In contrast to online RI. (Sutton & Barto 2018), the key challenge in offline RL is that the policy must be optimized solely based on the historical dataset D, without access to additional interactions with the environment.
2.2 TRANSFORMER BAANSFORMER-BAEDD OFFLINE RL METHODS
In recent years, the Transformer architecture (Vaswani et al. 2017) has been increasingly applied to offline reinforcement learning, primarily by treating a trajectory 7 { 0z, r2) 1 as a sequential input and representing its elements as unified tokens (Janner et al. 2021 Chen et al. 2021 Wu et al.)
2023 Wang et al. 2024). Specifically, the state 8t, action a,, and return G. G-=} اド "'r; are embed- ded into shared representation space, thereby enabling the model to capture long-range dependen- cies within the trajectory. Almost all Transformer-bast
offline RL methods, such as Transformer- based Behavior Cloning (BC) and Decision Transformer (DT) (Chen et al.] (02T), share a unified training objective, which can be formulated as:
K log mo (a Tgt))
L()) = -E
(1)
1=0
2

---

Published as a conference paper at ICLR 2026
where To denotes petiioommrr-prrmmteriizeeddd policy conditioned the trajectory prefix T<t (Ti)-o- In Transformer based BC, the prefix T<i consists of only (s,, 0,) pairs. While in DT, additionally incorporates the return-to-o token G, resulting in Tşx {(G,, 8;, a=)]== G, mainly serves as a conditioning signal that guides the policy towards trajectories with higher expected re- turns. For continuous action spaces, the loss in Eq. [ is equivalent to the mean squared error (Wang) & Bovik 2009), while for discrete action spaces, it corresponds to the cross-entropy loss (Mao ct | (023). This unified tokenization and sequential modeling paradigm allows Transformer-based offline RL methods to accommodate diverse task formulations within shared framework and has demonstrated strong modeling performance across various benchmark domains. Unlike traditional offline RL that depends on value-based or policy optimization objectives, Transformer-based meth- ods learn policies purely through sequence modeling.
METHODOLOGG
In this paper, we propose the Peak-Return Greedy Slicing (PRGS) framework, which comprises three fundamental components: (I) return estimator based Maximum Mean Discrepancy (MMD) that evaluates the potential return distribution for each state-action pair within a trajec- tory: (II) subtrajectory slicing strategy that greedily divides trajectories using an optimistically biased initial return estimate, with the resulting subtrajectorics used for policy training; and (III) an adaptive history truncation mechanism employed during policy evaluation and execution, which dy- namically determines whether to retain or discard historical information by comparing changes in the estimated value. The following sections provide a comprehensive description of each component.
3.1 MMD-BASED RETURN ESTIMATOR
To assess the potential value of each state-action pair within a trajectory at finer level of gran- ularity, we introduce an MMD-based return estimator. Maximum Mean Discrepancy (MMD) is a widely used metric for measuring the distance between distributions (Smola et al. (006). It quan- tifies the difference between two distributions X and Y in the Reproducing Kernel Hilbert Space (RKHS) (Berlinet & Thomas-Agnan 2004) and is defined as:
MMD (X, Y) X [k 2²] + Ey.y'>y (k(y, (22))-22m.. - y [k(z, v)],
where .) denotes a kernel function. In practice, MMD en- ables the estimation distributional differences using a finite set of sampled data, and has thus been extensively applied in tasks involving distribution fitting and alignment. In RL, value functions are often treated as distributions (Bellemare et al.] 2017 Dabney et al. 2018 Jullien al., 2025). since a single state-action pair can yield diverse future returns. Traditional quantile regression (Jurevckova] 2006) captures this distribu- tional nature by learning fixed set of quantile points. In con- trast, MMD offers a non- parametrie alternative for aligning distributions, allowing the return distribution of state-action pairs to be learned directly, rather than merely predicting an expected value. This approach captures uncertainty and the potential variability of state-action pairs more comprehensively, thereby improving the expressive- ness and robustness of value estimation.
Sort & Select Top-n
Estimation
Particles
Mean
MMD-based Estimator
(n)
Accurulatedd
Rá
Reward
8t
SE
0g
Figure 2: based optimistic return estimation.
Illustration of MMD-
Motivated by previous research (Nguyen- Tiguyen-TTang&a et al. 2021). we develop an MMD-based return es- timator. Given a single-step state s; and its corresponding action a, the estimator takes (st, a;) as input and produces set of N scalar samples:
Zo a.) {21, 22, 2N]; 2 R.
where ya denotes the model parameters. This set approximates the return distribution conditioned on the given state-action pair, rather than providing singular estimate. Moreover, unlike estimation methods that condition on entire trajectories, this design reduces the influence of irrelevant historical context. It focuses evaluating all possible intrinsic values of the current state-action pair.
3

---

Published as a conference paper at ICLR 2026
Ra
Peak
Casual Transformer
Oo
Og
02
0g
Selected Subtrajectories
Unused Subtrajectories
Mask
Casual Transformer
Letall
Og
02
Gx
Used Subtrajectories
Mask
Mask
Ris
Ri Peak
Casual Transformer
80
8
d0
01
O2
Slicing
Mask
Training
Figure 3: Illustration of the proposed greedy subtrajectory slicing mechanism. Given offline trajectory, compute the optimistic return Rt each time step and identify the peak point t as the split point for subtrajectory slicing. The selected subtrajectory is used for training while the remaining part is recursively sliced until all time steps covered.
To train the estimator, we minimize the MMD loss between the predicted distribution Zu ai) and the target distribution Ztage (st, ar)
{z1
, ;. shown in Eq
2):
1 N2
1
2
LMDD
k(z,, 23)
k(z (z),
(4)
N
N2
1=1
The target distribution Zurget computed based on the temporal difference target:
burge (5c, 0e) Ft + 7Z (st+1, as+1);
where wi denotes the target network with delayed updates. For simplicity, we adopt kernel function based the squared Euclidean distance (Eyster & White 1973), defined as k(z, y) -||x y/1². By minimizing the MMD loss above. the estimator encouraged learn a stable representation of the value distribution, which serves as reliable guidance signal for subsequent subtrajectory slicing.
3.2
GREEDY SUBTRAJECTORY SLICING
The greedy subtrajectory slicing module constitutes the core of PRGS. For clarity, provide de- tailed exposition of this module by dividing it into three sequential components: estimation, slicing, and training.
Estimation. Given an offline trajectory = {(sE, de, rz) length K, we commence by applying the MMD-based return estimator introduced in the previous section to estimate the fu- ture return distribution for each state-action pair (st, ax)) along the trajectory. To enable consistent comparison across different timesteps, we then derive scalar scores with optimistic bias from the estimated return distributions. According to Eg. the estimator generates N particles at each timestep representing an approximate distribution of returns starting from (s). To derive an optimistic scalar estimate, the particles are sorted in descending order:
(1,(1) 2 22) 2 2 (,)N)-
The value function approximation is then defined as the mean of the top-n particles:
(s.. a.)
21(0)
(5)
n
where E {L...., N} is a tunable hyperparameter. Smaller values of n yield estimates closer to the upper quantiles of the distribution, resulting in more optimistic evaluation, while larger values lead conservative estimates. Intuitively, ((n) captures the potential upper-bound
4

---

Published as a conference paper at ICLR 2026
return by focusing the most favorable outcomes, thereby preventing suboptimal realizations from disproportionately degrading the evaluation of specifie subtrajectories.
According to C.2). @(ah) denotes an optimistic estimate of the return starting from timestep t. To ensure comparability across different timesteps, we map this quantity to the perspective of the trajectory's 8 starting point (8o, an) by incorporating the cumulative realized reward obtained before t, as shown in Figure Formally, we define the aligned optimistic return from the start as:
RJ Σ
Yn + Q.Te
In this way, it' denotes the optimistic total return the agent would receive by following the observed trajectory up to step and then executed action de at state 8t. This alignment brings the potential returns at different timesteps onto a unified scale, enabling consistent comparison and supporting the subsequent greedy segmentation of subtrajectories.
Slicing. For the definition of Rg if there exists timestep t" [0, K] such that R
maxosisk Rij then the subtrajectory To.- {(sp, ap, ro)...., (s:,, ar" re-) } identified as the optimal subtrajectory within the trajectory. In other words, learning from this segment is expected to yield nearly optimal return. Any trajectory segment occurring after the timestep t* degrades the estimation of ft. resulting a suboptimal return value. Based on this observation, we implement greedy slicing strategy: for each complete trajectory. we select Th. corresponding to the peak optimistic return Ry and use it as training sample for the Transformer model. Accordingly, t* serves as the split point.
Training. During training, only the timesteps within To:- retained, while the other steps masked and excluded from the loss computation. Following Eq.II) we obtain the modified loss function as:
(0) = EEWD Σ log (e Toa-),
(=0
where O and t" denote the starting and ending timesteps of the selected subtrajectory, respectively.
Then, the selected subtrajectory Tou- is labeled as used, and the remaining unused portion TiK- ((51+4, 1 ar +1, 1,, (1, (SK,, (K,TK)) is recursively segmented using the same greedy strategy. In the m-th recursion step, a new split point is identified to extract high-quality subtrajectory. and the corresponding loss value 2. calculated. This process continues until all timesteps are covered, resulting in complete decomposition of the original trajectory into a set of disjoint subtrajectories, as illustrated in Figure 3 The overall optimization objective for a training iteration is defined as the weighted sum of the losses across all extracted subtrajectories:
M
Liotal
Σ Am
where M denotes the number of resulting subtrajectories, which dynamically varies depending on the specific input trajectory. E [0, 1] denotes a weighting coefficient, indicating that the contri- bution of L. diminishes as m increases. In particular, when À = O, only the first subtrajectory is considered, whereas À = indicates that all subtrajectories are treated uniformly during training. This design aligns with the objective of guiding PRGS to focus more on subtrajectories associated with higher values of Rg. Meanwhile, this recursive slicing framework ensures that every timestep contributes to training.
3.3 ADAPTIVE HISTORY TRUNCATION IN EVALUATION
During the training phase, PRGS applies greedy segmentation to divide the entire trajectory into multiple subtrajectories and trains on each independently. This design makes the initial state of each subtrajectory critical for subsequent decision-making while the preceding historical context is not explicitly incorporated into training. Consequently, unconditionally retaining all historical information during evaluation introduces a mismatch between training and evaluation, potentially degrading policy performance.
5

---

Published as a conference paper at ICLR 2026
To mitigate this inconsistency, PRGS adopts an adaptive history truncation mechanism during eval- uation. At each timestep t (t 2 1), the return estimator provides an optimistic evaluation of the current state, which is then compared with the estimate from the previous step to determine whether the historical trajectory should be retained. First, we define the state value estimate at time step as:
vs(ss) (5-1((--1.
1,
where Q,-0 denotes an optimistic return based on Eq. In the evaluation phase, PRGS compares the estimated value of the current state against that of the preceding step: AV, = V((5;) V (s If AV, > O, it suggests that the current state holds greater potential for return than the previous one, implying that the preceding trajectory no longer contributes valuable infor- mation for future decisions. In such cases, PRGS discards the historical trajectory and retains only the current state as the new starting point, from which subsequent decisions are derived. Formally, the history update mechanism during evaluation can be written as:
H2 (st).
if AV, > 0,
(6)
U (sr), otherwise,
where H, denotes the retained trajectory at timestep t. This adaptive history truncation strategy ensures consistency between training and evaluation. While training relies on subtrajectories start- ing from intermediate states, the evaluation phase can also dynamically reset the history based on the return estimator, thereby mitigating the impact of irrlevant or low-quality past information on current prciision-maing
In summary, the PRGS framework systematically enhances the ability of Transformer-based of- fline RL through three tightly integrated modules. The MMD-based return estimator provides op- timistic value estimates for single state-action pairs, while greedy slicing identifies and prioritizes high-return subtrajectories for training. During evaluation, adaptive truncation dynamically discards irrelevant history to align inference with training-time segmentation. Together, these components enable timestep-level subtrajectory selection and composition, offering a concise and effective way to improve the stitching ability of Transformer-based offline RL. The implementation details and algorithmic description of PRGS can be found in Appendix
4 RELATED WORK
4.1 OFFLINE REINFORCEMENT LEARNING
Besides the Transformer-basee approaches discussed in this paper, several other representative meth- ods have also been developed in onf RL. One prominent line of work involves policy constraint methods, such as BCQ (Fujimoto et al. 2019 and BRAC (Wu et al., 2019), which mitigate the bias introduced by out-of -itoof--isttibutionn (OOD) Mao et al., 2024b) actions by constraining the learned policy to remain close to the behavior policy during updates. Another class of methods focuses on advantage weighting and value function regression. For instance, CQL (Kumar et al. 2020 imposes a conservative penalty on the O-function to prevent overestimation of OOD actions. At the same time, IQL (Kostrikov et al. (002)) leverages expectile regression to stabilize the learning dynamics. With the development of sequence modeling techniques, researchers have introduced Diffusion Models (Ho et al. 2020)
more scalable tool for policy representation and planning. Typical representatives include Diffuser (Janner et al. 2022) and DD (Ajay et al. 2023). In addition, EDP (Kang et al. (2023) and DiffuserLite (Dong al. 2024 optimize the training and inference effi- ciency of diffusion-based policies. This series of works demonstrates the great potential of diffusion architectures in Offline RL, bringing a new generative model paradigm to offline decision-makinn
4.2 STITCHING IN TRANSFORMER BASED OFFLINE RL METHODS
Although Transformer-based approaches have shown strong modeling capabilities in offline RL, they still face limitations in stitching together high-value segments from different parts of trajecto- ries. Several recent works have explored ways to improve this ability. For instance, TT (Janner et al.) 2021) incorporates discretization and dynamic programming into trajectory modeling, enabling the model to utilize planning mechanisms for more effective trajectory recombination. Methods such

---

Published as a conference paper at ICLR 2026
as RCSL (Srivastava et al. [019| Schmidhuber 2019), QDT (Yamagata al. 2023) CGDT (Wang et al. 2. and QT (Hu et al.
2024) integrate value estimation into sequential modeling, allowing value- guluedd composition of trajectory segments. EDT (Wu et all 2023) estimates the value of past observations and adaptively determines the effective context length based on the estimated utility. Meanwhile, DoC (Yang et al., (023) takes alternative approach by learning latent representations of future trajectories and conditioning the policy model on these representations to guide behavior Wu (2023) propose using reward machines structure offline data, enabling strong policies to be learned from less but higher-quality data in complex tasks.
adjustment. And And[Sun [Sun
The above methods have demonstrated effectiveness in improving trajectory stitching. Nevertheless, they lack a fine- Sine-griined perspective to explicitly distinguish between high-quality and low-quality subtrajectories within trajectory. Most approaches operate at the trajectory level or rely on implicit latent representations, without mechanisms to explicitly identify high-quality subtrajectories at the timestep level. To mitigate this constraint, we present an interpretable mechanism that identifies split points and directly selects informative subtrajectories from given trajectory. It is compatible with various Transformer-basee offline RL algorithms and significantly improves their stitching performance.
5 EXPERIMENT
This section empirically evaluates the effectiveness and advantages of PRGS through a series of sys- tematic experiments. The evaluation includes four main components: (I) integrating PRGS into sev- eral representative Transformer-basee offline RL algorithms and assessing its overall performance across several benchmark datasets; (II) performing ablation studies to investigate the contributions of individual components and hyperparameters, including the number of particles n in the MMD- based estimator and the adaptive history pruning mechanism; (III) verifying the efficacy segmen- tation at the timestep level by comparing it with conventional trajectory-level filtering methods; (IV) finally, providing visualizations to illustrate the inner workings of PRGS to facilitate understanding of its impact on the training.
5.1 OVERALL PERFORMANCE
We conduct experiments on various representative offline RL benchmarks thoroughly assess the effectiveness of PRGS. The baselines are grouped into three categories: (1) classical offline RL al-
gorithms, including CQL (Kumar et al. [2020) IQL (Kostrikov al., 2022). BEAR (Kumar al.)
2019). and TD3+BC (Fujimoto & Gu 2021); (2) recent approaches emphasizing stitching capabil- ity, such as QDT (Yamagata et al. 2023). CGDT (Wang et al. 2024), and EDT (Wu et al. 2023); and (3) Transformer- framework with enhanced representational capacity while still adhering to Eq. [ and is therefore adopted as an important baseline in this paper. To demonstrate the effectiveness of the proposed framework, we construct Transformer-based offline RL variants augmented with PRGS, denoted by appending PRGS as a suffix. The evaluation spans diverse domains, including standard continuous control tasks from the D4RL benchmark (Fu et al. 2020), natural language and multi-step reason- ing tasks in Ba Chevalier-Boisverr NuctionNet (Su et al. 2024). These benchmarks cover both continuous and discrete action spaces, allowing for more comprehensive comparison. For stitching-focused methods, we restrict the eval- uation to scenarios where such capabilities are particularly relevant. In the experimental results, the arrows denote the performance difference between the PRGS algorithm and the baseline algorithm. All experiments are based on publicly available and reproducible implementations, with detailed hyperparameter settings and dataset configurations provided in Appendix A and Furthermore, additional experimental results and a more detailed analysis are presented in Appendix C
methods like BC, DT, and PDIT (Mao et al. 2024a) PDiT establishes
al. 2019), and auction-style decision-making tasks in
The experimental results on the D4RL benchmark (see Table D show that the proposed PRGS framework consistently achieves superior performance across a wide range of tasks. For Gym tasks, traditional offline RL algorithms such as CQL and IQL perform reasonably well on medium or expert datasets, but exhibit a notable drop in performance ods enhanced with PRGS demonstrate significantly improved stability and higher average returns. For instance, DTRP shows substantial gains over vanilla DT on mediun-replay datasets. In the Adroit and Kitchen domains, PRGS excels on several challenging human demonstration datasets (including pen-human and hammer-human), outperforming existing baselines and demonstrating
medium- datasets. In contrast, meth-
7

---

Published as a conference paper at ICLR 2026
Table 1:
Offline RL methods on D4RL benchmarks.
Scenarios
COL
JQL
BEAR
TD3+BC
BC
BC-PRGS
DT
DE-PRGS
PDIT
PDIT-PRGS
hulfeheetah-mediiee meduuee--emiieeerrrrr
91.6
86.7
53.4
90.7
96.2+9.4
88.42.9
01.7+0.3
93.1+04
73.0+4.3
86.1=288
hopper-mooding mall----iiiiu--ppert
105.4
91.5
96.3
98
07.5=13.1
71.5+0000
8:0.5
111.2:0.3
410.1
111.3:0.1
saller2d-mediin-exxpp
108.8
100.6
40.1
110.1
109.7+0.3
108.5+0.3
1089++0.1
109.6+0.4
106.
106.6-1.8
halfcheetah medium
49.2
47.4
41.7
48.4
40.520.1
40.720.3
40.0::011
41.820.5
42.8423
39.9=00.1
Gym
hopper-
09.4
06.3
02.1
50.3
10.5+115
62.46577
63.6+266
93.1+2.7
68.242.44
92.4-2.0
Tasks
walker2d mediuum
83.0
78.3
59.1
83.7
78.84:11
77.842.
78. 1.
78.520.9
77.64006
77.9=2.2
45.5
44.2
38.6
44.8
35.8+0.7
37.9++0.7
35.0+1.0
39.9+0.3
40.8+22.3
34.5-0.4
aapper-medium-reppe
55.0
94.7
33.7
00.9
48.0=228.2 2
30.5+23.0
78.74:0.3
98.1
80.6&22.7
96.4=116
walker2d mediimm replay
77.2
73.9
19.2
81,8
57.543.3
70.8441.3
71.51177
81.120.8
71.110.6
78.2-4.11
Average
80.0
77.0
08.2
75.3
01.8
67.7 122
75.3
62.9 m7
76.3
80.1 133
pen- human
37.5
71.5
-1.0
-3.9
109.917.6
118 118.112..8
63.6.17.2
118.916.9
97. 97.713999
116.61:3.0
Tummer-human
4.4
1.4
0.3
1.0
2.5+0.8
7.2++4.3
3.0+0..8
4.4+220
3.7+1.7
6.8+33.6
Adroit
door-human
4.3
0.3
14.5+119
13.5+3.3
22 22.7444..4
7,8+1.6
14.1+-4.6
8.9+5.3
Tasks
pen-cloned
39.2
37.33
96.5
5.1
77.8+12.3
60.1+5.9
60.3++8.7
60.++633
91.4++1..0
80.6+16..
summer-cloned
2.1
0.3
0.3
2.6+224
2.0+1.77
1.6+0.4
4.3+2.2
3.8+1.9
6.7=0.0
doot-clonod
0.4
1.6
-0.1
-0.3
1.6+1.0 00
0.3+0.3
14.3+1177
5.3+1.88 1.88
9.5++2.3
7.4+0.5
Average
15.6
19.7
4.3
0.3
318
35.1
30.9
35.2
36.7
37.7
kichets-complete
43.8
62.5
0.0
0.0
48.3+338
70.8++2.9
40.3+5.33
71.7+-2.9
27.5++0.2
31.7-633
Kitchen
Taska
kitcher-portial
49.8
46.3
0.0
52.5112.5
59.9:1::2
50.2188.0
43.3188.0
50.8:2.9
Average
46.8
54.4
6.6
0.0
50.4
60.4 1080
30.1
65.5 1154
35.4
41.3 188
mareld- 2a-aaareee
42.1
14.8
16.42-4.7
33,9211.5
60.317777
12.41:1
73:2:11.6
82.3:5.0
Maze2D
Tasks
maze2d-me mazeld-mediam
41.8
34.9
26.0
62.1
20.1+844
10.3+1.8
37.0++.3
90.4+6.8
51.2+44.9
86.6+15.2
maze2d -arga
49.6
61.7
81.0
88.8
10.34888
9.3
26.43-4.9
127.5
40:0=:10:2
106.6=20.4
Average
62.0
46.2
57.2
55.2
15.6
20.8 151
40.9
100.1 1992
54.8
91.8 157
antmaze- umaaze
74.0
87.5
78.8
$2.32:322
67.0.117
69.3114.9
71.11130
AntMaze
antmaxe umauee dinieeeeeeeee
84.0
62:2
61.0
71.4
81.7199.5
83.011..
63.011:.66
80.113.8
60.31:0.1
00.3:9.00
Tasks
antmare-mediumm-iivrs
70.0
3.0
2.0+-1.0
2.7+1.2
1.3+0.6
3.3+5.8
0.0+0.0
antmaze- large-diveree diverse
14.9
47.5
0.0
0.0
0.0:10.0
0.0:10.0
1.011.0
16.717.8
1.7.11.5
11.013.0
Average
66.8
57.2
38.3
41.5
42.2 107
33.4
115.
strong robustness to sparse rewards and noisy trajectories. On Maze2D and AntMaze tasks, which involve long-horizon planning and trajectory stitching, PRGS again shows effectiveness. Notably, DT-PRGS achieves a score of 127.5 maze2d-large, outperforming all compared methods. Over- all, the PRGS family consistently matches or exceeds the best best--perfrming baselines across diverse settings. These findings validate the effectiveness of the proposed explicit subtrajectory segmenta- tion mechanism in improving the performance of Transformer-baar based offline RI.
The results on Gym medium and medium-replay tasks are reported in Table 2 Overall, the DT-baaed PRGS method consistently outperforms other variants. In particular, on the coopper-meduu--rpllee and
tasks, PRGS achieves significantly better performance than CGDT, demonstrating its ability to better capture critical subtrajectories in environments with complex dy- namics and long-tern term dependencies. On average, PRGS ranks first among the five methods, achiev- ing a 10.9 10.9-point improvement over vanilla DT. These results indicate that PRGS effectively enhances performance, especially excelling in scenarios where replay data is noisy and cross-trajectory recom- position is required.
To highlight the practical utility of the proposed algorithm, we further evaluate various baselines and their PRGS-enhanced variants on the AuctionNet benchmark. As shown in Table 3 PRGS consistently improves the performance scores of the original algorithms in the advertising bidding scenario, with particularly notable gains for BC, This may be attributed to the relatively simple mechanism of BC, which offers greater potential for performance enhancement on AuctionNet.
Finally, in the BabyAl benchmark, evaluate multiple tasks of varying difficulties. The PDiT method, which employs a multi-layer Transformer architecture, achieves strong performance in this domain. rrepots the results of several Transformer-baat based offline RL methods alongside their PRGS variants. Across almost all tasks, the PRGS variants consistently outperform the original algorithms, with particularly notable gains the more challenging tasks.
5.2 ABLATION STUDIES
To further disentangle the contribution of each proposed mechanism, extensive ablation studies were conducted. The particle number n. was first examined assess the influence of optimistically biased
8

---

Published as a conference paper at ICLR 2026
Table 2: Results on Gym medium and medium-replay tasks.
Table 3: Performance comparison in different periods on AuctionNet.
Gyms Tasks
DT
QUT
EDT
CGDT
PRGS
DE-PRRSS
PINT
balfchheetah-ccedras
hoppan-mmodiam
63.6
40.0
42.2
62.5
96.9
43.0
03.1+2.7
41.82:l.5
P14
P15
P16
294.9123114
230 01..82117.4
855
100.1:125
174.
300 13.9
282 292.0:33.44
200.2112
270.288.9
297.0
278.5
250.4-21.1
285.1414.77
301,
254.4.225.1
311-442211.0
265.2..7.11
30L11290
214.22300.5
walker2d medium
78.1
70.1
72.8
79.1
78.5..0.0
hattcheetah seedram replay
83.0
35.7
37.8
40.4
32.2200.
PI8
210.3+100
20-1
216.5+11.1
233.4+19.2
23666-154
78.7
96.1++1.7
P19
286.4::16.00
202.2=16.2
285.1+338.3
264.9=12.7
9:11.
285.2+122.7
walker?d medium replay
71.3
59.1
74.8
78.1
$1.1100.
P30
241 311.66111.6
24
241.0:55
232 12.6
241 21:11.
Average
6122
54.6
63.4
71,8
72.1 -00
Average
267.3
271.9
Table 4: Results on BabyAl tasks.
DT
PRGS
PRGS w/o AHT
Performance
Scenarios
BC
BC-PRG
DT
DT-PRG
PDTT
POCE-PRGS
200
GoToRedBall
87.9+1.7
80.0+33
00.7+222
9166+15
99.240.1
98.5+0-1
GoToLocal
77.2+0.3
63.066.0
81.441.7
82.6+35
91.41188
95.1+0044
50
GoToSey
31.4+:4.0
38.7+1.3
50.0+110
35.5t:1.1
$8.7+5.4
PutNeutlooal
17.5+-5.6
19.0+1.7
13.1+16
12.1+4.7
23.9+2.9
32.7+5.1
ColoetLocalDist
50 113.6
50.51336
597119
61.8...7
72.31:4.7
32 Number of particles (n)
BossLevel
36.32332
42.2322.
420011.4
43.2155.3
49.0.1-4.3
48.0133.
Average
32.1
55.3 152
54.3
65.9 1L.
65.11
65.9
Figure Ablation studies.
estimates, and the necessity of the adaptive history truncation mechanism as an alignment strategy during evaluation was also investigated.
Experiments on the Maze2D dataset provide the ablation results summarized in Figure 4 When the particle number is very small, the results are suboptimal, likely because the MMD-based estimator generates outlier estimates that impair accuracy. When excessively large, the method degenerates into a standard value estimator, losing its optimistic bias and resulting in degraded performance. On the other hand, the variant denoted as PRGS w/o AHT, which excludes the adaptive history truncation mechanism, shows a clear performance drop, indicating that the adaptive history truncation mechanism plays a crucial role maintaining alignment between training and evaluation. Additional ablation results about the coefficient À are reported in Appendix C
5.3 EFFECTIVENESS OF TIMESTEP LEVEL SUBTRAJECTORY SLICING
To validate that the performance improvement achieved by PRGS arises from its explicit timestep- level subdivision rather than merely filtering trajectories for training, we evaluate several PRGS variants and compare them with results obtained from training Maze2D datasets filtered to retain only a small fraction of high-quality trajectories.
Vanillaa
PRGS
Experimental results in Figure 5 demonstrate that timestep- level subtrajectory slicing significantly enhances the perfor- mance of various baseline methods. In particular, models trained on the top 10% or 20% of trajectories do not neces-
Purformance
100
sarily outperform the original algorithms, a phenomenon es- pecially evident in BC. By contrast, PRGS consistently sur- passes other slicing strategies across all three base models, confirming the effectiveness and generality of timestep-level subtrajectory slicing in improving policy learning. These findings suggest that finer-grained subtra- jectory selection enables more stable and efficient training across diverse offline RI. frameworks.
DT
PDIT
Methods
Figure 5: Results of different trajec- tory selection strategies.
5.4 VISUALIZATION
In this section, we present visual analysis of trajectory data in maze2d-meding Figure O shows that the raw trajectories span most of the maze, yet the majority fail to reach or approach the target (marked by red star). In Figure trajectory points are color-coded by their return-to-go (Rtg), with brighter colors corresponding higher cumulative returns; this highlights that high-return re- gions are concentrated around the target. Figure ((c) illustrates the estimates of Q for state-action pairs obtained from the MMD-based estimator. While the overall trend aligns with the distribution of return-to-go additionally distinguishes value variations in distant regions, suggesting that the method captures local values more effectively and yields more optimistic estimates. Finally, Fig- ure d) depicts the first subtrajectory selected by PRGS during the slicing process. PRGS consis- tently favors subtrajectories that include high-value regions while discarding low-quality segments

---

Published as a conference paper at ICLR 2026
that drift away from the target, thereby showing its ability automatically extract the most infor- mative portions of complex trajectory sets for learning.
CSSS (b) Rig visualization.
(a) Trajectories.
(c) Q estimation.
(d) Subtrajectories.
Figure 6: Visualization of trajectory data in maze2d-medium.
CONCLUSION
This paper introduces Peak-Return Greedy Slicing (PRGS), which integrates an MMD-based return estimator, greedy subtrajectory slicing strategy, and adaptive history truncation mechanism to explicitly select high-quality subtrajectories at the timestep level. This approach substantially improves the compositional capability of transformer-basee offline RL. Experimental results multiple benchmark tasks validate the effectiveness of PRGS, highlighting its potential as a concise and efficient framework and suggesting promising directions for extensions environments. And extending this work to non-transformer- based methods is interesting direction.
complex
REPRODUCIBILITY STATEMENT
We make every effort to ensure reproducibility of our work. The paper provides a complete description of the proposed PRGS framework in Section including algorithmic details and train- ing objectives. All experimental settings, benchmark descriptions, and evaluation metrics are re- ported in Section 5 with additional hyperparameter details and implementation notes included in Appendix A The datasets employed (D4RL., AuctionNet, BabyAl) are publicly accessible, and describe the preprocessing steps in Appendix B An link to the source code is provided https: /ttp:////tttttt com/ de iigett fool /PRGS and the code is also available in the supple- mentary material together with instructions for running the experiments.
ACKNOWLEDGMENTS
This work was supported by the National Natural Science Foundation of China (Grant No.62506210).
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh Interna- tional Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OvenRevie 2023.
Mare G. Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement learning. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Con- ference Machine Learning, ICML Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 449 458. PMLR, 2017.
Alain Berlinet and Christine Thomas- Agnan. Reproducing kernel hilbert spaces probability and statistics. 2004.
10

---

Published as a conference paper at ICLR 2026
David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return- couur--ondiiioeed supervised learning work for offline reinforcement learning? In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Sys- tems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, viriual, pp. app15000 15084-15097, 2021.
Maxime Chevalier-Boisvert. Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Will Dabney, Mark Rowland, Mare G. Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Sheila A. Mcllraith and Kilian Q. Weinherger (eds.), Pro- ceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Ed- ucational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 2892-2901.. AAAI Press, 2018. doi:
Daniel C Dennett. The intentional stance. MIT press, 1989.
Zibin Dong, Jianye Hao, Yifu Yuan, Fei Ni, Yitian Wang, Pengyi Li, and Yan Zheng. Diffuserlite: Towards real-time diffusion planning. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neural In- formation Processing Systems 38: Annual Conference on Neural Information P Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10- 15, 2024, 2024.
James W Byster and JA White. Some properties of the squared euclidean distance location probleml. AIIE Transactions, 5(3):275-280, 1973.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, htss/2044.02199, 2020.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6- 14, 2021, virtual, pp- 20132-20145, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, Cali- fornia, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2019.
PMLR,
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare" Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
33: Annual Conference on Neural
Shengchao Hu, Ziqing Fan, Chaoqin Huang, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Q-value regularized transformer for offline reinforcement learning. In Forty-first Interna- tional Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. Open- Review.net, 2024.
11

---

Published as a conference paper at ICLR 2026
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 1273-1286, 2021.
Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. Kamalika Chaudhuri, Stefanie Jegelka, Le Song. Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 9902-99155 PMLR, 2022.
Sami Jullien, Romain Deffayet, Gaan-Michel Renders, Paul Groth, and Maarten de Rijke. Distribu- tional reinforcement learning with dual sspectile-quantie regression. In Silvia Chiappa and Sara Magliacane (eds.), Conference on Unceriuinty in Artificial Intelligence, Rio Othon Palace, Rio de Janeiro, Bruzil, 21 1-2 July 2025, volume 286 of Proceedings of Machine Leurning Research, pp- 1909-1923. PMLR, 2025.
Jana Jurevckova. Quantile regression. Journal of the American Statistical Association, 101:1723- - 1723, 2006.
Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Sungyoon Kim, Yunseon Choi, Daiki E. Matsunaga, and Kee-Bung Kim. Stitching sub-trajectories with conditional diffusion model for goal-conditioned offline RI. In Michael J. Wooldridge, Jennifer G. Dy, and Sriraam Natarajan (eds.), Thirty-Eighth AAAI Conference Artificial In- telligence, AAAI 2024, Thirty-Sixth Conference on Innovative Applications of Artificial Intel- ligence, IAAI 2024, Fourteenth Symposium on Educational Advances in Artificial Intelligence, EAAI 2014, February 20-27, 2024, Vancouver, Canada, pp. 13160-13167. AAAI Press, 2024. doi: 10.1609/AAAI V38112.2291..
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché Buc. Emily B. Fox. and Roman Garnett (eds.), Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 20/9, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11761-11771, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Ties Lin (eds.), Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Aviral Kumar, Anikait Singh, Frederik D. Ebert, Mitsuhiko Nakamoto, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre-training for robots: Offline RL enables learning new tasks in a handful of trials. In Kostas E. Bekris, Kris Hauser, Sylvia L. Herbert, and Jingjin Yu (eds.), Roboties: Science and Systems XIX, Daegu, Republic of Korea, July 10-14, 2023, 2023. doi: 10.15607/ D.SS2203.XXXX.C
Jaewoo Lee, Sujin Yun, Taeyoung Yun, and Jinkyoo Park. GTA: generative trajectory augmentation with guidance for offline reinforcement learning. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Process- ing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
12

---

Published as a conference paper at ICLR 2026
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. CoRR, abs/2005.01643, 2020.
Guanghe Li, Yixiang Shan, Zhenghang Zhu, Ting Long, and Weinan Zhang. Diffstitch: Boosting offline reinforcement learning with diffusion-baset trajectory stitching. In Forty-first Interna- tional Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. Open- Review.net, 2024.
Haohong Lin, Wenhao Ding, Zuxin Liu, Yaru Niu, Jiacheng Zhu, Yuming Niu, and Ding Zhao. Safety-aware causal representation for trustworthy offline reinforcement learning autonomous driving. IEEE Robotics Autom. Lett., 95):4639-4646, 2024. doi: 10. 10/.11999...2243333909
Anqi Mao, Mehryar Mohri, and Yutao Zhong. Cross-entropy loss functions: Theoretical analysis and applications. ArXiv, abs/2304.0722 007288 2023.
Hangyu Mao, Rui Zhao, Ziyue Li, Zhiwei Xu, Hao Chen. Yiqun Chen, Bin Zhang, Zhen Xiao, Junge Zhang, and Jiangjin Yin. Pdit: Interleaving perception and decision-makin transformers for deep reinforcement learning. In Mehdi Dastani, Jaime Simão Sichman, Natasha Alechina, and Virginia Dignum (eds.), Proceedings of the 23rd International Conference on Autonomous Agents and Multiagent Systems, AAMAS 2024, Auckland, New Zealand, May 6-10, 2024, pp. 1363-1371. International Foundation for Autonomous Agents and Multiagent Systems / ACM, 2024a. doi: 10.5555/363567.3662999
Yixiu Mao, Qi Wang, Chen Chen, Yun Qu, and Xiangyang Ji. Offline reinforcement learning with OOD state correction and OOD action suppression. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, UIrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Conference Neural Information Process- ing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024b.
Ryo Nakahashi, Chris L. Baker, and Joshua B. Tenenbaum. Modeling human understanding of complex intentional action with bayesian nonparametrie subgoal model. In Dale Schuurmans and Michael P. Wellman (eds.), Proceedings of the Thirtieth AAAI Conference on Artificial Intel- ligence, February 12-17, 2016, Phoenix, Arizona, USA, pp. 3754-3760. AAAI Press, 2016. doi: 10.1609/AAAL V3011.9889
Thanh Nguyen-Tanne Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning via moment matching. In Thirty- Fifth AAAI Conference Artificial Intelligence, AAAI 2021, Thirty- Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, Febru- AAAI Press, 2021. doi: 10.16099/AAAL.333110001
ary 2-9, 2021, pp.
Yue Pei, Hongming Zhang, Chao Gao, Martin Müller, Mengxiao Zhu, Hao Sheng, Haogang Zhu, and Liang Lin. Double check my desired return: Transformer with target alignment for offline reinforcement learning. arXiv preprint 2508 2025.
Adena Schachner and Susan Carey. Reasoning about 'irrational' acttions: When intentional move- ments cannot be explained, the movements themselves are seen as the goal. Cognition, 129(2): 309-327, 2013.
Roger Schank and Robert P Abelson. Scripts, plans, goals, and understanding: An inquiry into human knowledge structures. Psychology press, 2013.
Jürgen Schmidhuber. Reinforcement learning upside down: Don't predict rewards just map them to actions. CoRR, abs/1912.02877 2019.
Alexander J Smola, A Gretton, and K Borgwardt. Maximum mean discrepancy. In 13th international conference, ICONIP, volume 6, 2006.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmidhu- ber. Training agents using upside-down reinforcement learning. CoRR, abs/1912.02877 2019.
13

---

Published as a conference paper at ICLR 2026
Kefan Su, Yusen Huo. Zhilin Zhang, Shuai Dou, Chuan Yu, Jian Xu, Zongqing Lu, and Bo Zheng. Auctionnet: A novel benchmark for decision-making in large- aare--ccale games. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NewrIPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
Haoyuan Sun and Feng Wu. Less is more: Refining datasets for offline reinforcement learning with reward machines. In Noa Agmon, Bo An, Alessandro Ricci, and William Yeoh (eds.). Pro- ceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS 2023, London, United Kingdom, 29 May 2023 2 June 2023, pp. 1239-1249 ACM, 2023. doi: 05555535555946.35888666
Richard S. Sutton and Andrew G. Barto. Reinforcement learning an introduction, 2nd Edition. MIT Press, 2018
Laurens van der Maaten and Geoffrey E. Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9:2579-2605, 2008.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 9998-6008, 2017.
Yuanfu Wang, Chao Yang, Ying Wen, Yu Liu, and Yu Qiao. Critic-guided decision transformer for offline reinforcement learning. In Michael J. Wooldridge, Jennifer G. Dy, and Sriraam Natarajan (eds.), Thirty-Eight AAAI Conference on Artificial Intelligence, AAAI 2024, Thirty-Sixthe Confer- ence on Innovative Applications of Artificial Intelligence, IAAI 2024, Fourteenth Symposium on Educational Advances in Artificial Intelligence, EAAI 2014, February 20-27, 2024, Vancouver, Canada, pp. 15706-15714. AAAI Press, 2024. doi: (0.1609/AAAALV338144422022
Zhou Wang and Alan Conrad Bovik. Mean squared error: Love it or leave it? a new look at signal fidelity measures. IEEE Signal Processing Magazine, 26:98-117, 2009.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. CoRR, abs/1911.11361, 2019.
Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Anmual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023.
Xin Xin, Tiago Pimentel, Alexandros Karatzoglou, Pengjie Ren, Konstantina Christakopoulou, and Zhaochun Ren. Rethinking reinforcement learning for recommendation: A prompt perspective. In Enrique Amigó, Pablo Castells, Julio Gonzalo, Ben Carterette, J. Shane Culpepper, and Gabriella Kazai (eds.), SIGIR 22: The 45th International ACM SIGIR Conference on Research and Devel- opment in Information Retrieval, Madrid, Spain, July 11 15, 2022, pp. 1347-1357. ACM, 2022. doi: 10. 10.1145/347749553333777
Jian Xu, Zhilin Zhang, Zongging Lu, Xiaotie Deng, Michael P Wellman, Chuan Yu, Shuai Dou, Yusen Huo, Zhiwei Xu, Zhijian Duan, et al. Auto-bidding in large-scale auctions: Learning decision-making in uncertain and competitive games. In NeurIPS 2024 Competition Track, 2024.
Taku Yamagata, Ahmed Khalil, and Raúl Santos Rodrigue. Q-learning decision transformer:
Leveraging dynamic programming for conditional sequence modelling in offline RL. In An- dreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt. Sivan Sabato, and Jonathan Scarlett (eds.), International Conference Machine Learning, ICML 2023. 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 38989- 39007. PMLR, 2023.
14

---

Published as a conference paper at ICLR 2026
Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Separat- ing what you can control from what you cannot. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Ziqi Zhang, Jingzehua Xu, Zifeng Zhuang, Jinxin Liu, and Donglin Wang. Context-former: Stitch- ing via latent conditioned sequence modeling. CoRR, abs/2401. 16452, 2024. doi: 10.48550/
4RXIV.2401.16552
15

---

Published as a conference paper at ICLR 2026
A IMPLEMENTATION DETAILS
A.1 ALGORITHMIC DESCRIPTION
The algorithm for PRGS summarized in Algorithm supplementary material.
The code for PRGS be found in the
Algorithm 1: Peak-Return Greedy Slicing (PRGS)
Input: Offiin dataset D {7): Return estimator Zu; Target estimator Z
Transformer-baat policy ms; Discount factor y, Number of particles N; Number of Chosen top particles nı Coefficient X
Training Phase:
for each trajectory {(s, at, rrl E D do Sample particles Z and target particles Zuaget
Compute MMD loss and update parameters vi of the estimator
Periodically update target estimator 7 + Z Initialize unused timestep set U {0, ....... K} while not empty do
Litaa ppeev +-
10
m 4-1 1
11
for E 2 do
Obtain N particles: (st, 0z) = {2.1, 21,N 21,N
13
Sort in descending order to get 21,(1) ... ...2..0 (N(NN)
Compute optimistic value: @(()) 1 - (t.
"
Compute aligned optimistic return: Ry + I + (m)
arg max RJ tell Select subtrajectory :tp and mask all other steps
16
17
Compute loss L. via masked Transformer forward pass Mark (
19
as used and set 4 +- (o.....
20
ppev + -t +1
Eioall + Ciotal + A
21
22
m f-
23
Update the parameters 0 of the policy
24
Evaluation Phase:
25
for each rollout episode do Initialize history H <- (
26
for each step do if then
28
Estimate value V1(8t) < ②(a) (8-1190))) Va:)) > ( then Discard history: H {(t)
30
Predict next action no (H) Append (s,, a,) to
33
A.2 HYPERPARAMETERS
Hyperparameters were based on the implementation of Decision Transformer and are listed in Ta- ble 5 Most shared hyperparameters were used without extensive tuning. All experiments in this paper are run on Nvidia GeForce RTX 3090 graphics cards and Intel(R) Xeon(R) Platinum 8280 CPU.
16

---

Published as a conference paper at ICLR 2026
Table 5: Hyperparameters used in experiments.
Hyperparameter
Value
Number of layers
3
Number of attention heads
I
Embedding dimension
128
Activation function
ReLU
Positional encoding
No for D4RL and BabyAl, Yes for AuctionNet
Batch size Context length
64
20
Dropout
0.1
Learning rate
le-4
Grad norm clip
0.25
Weight decay
le-4
Learning rate warmup
linear warmup for 10000 training steps
Training epochs Eval episodes
10
32
Discount factor 7
1
Number of particles N
128
Number of Chosen top particles n
8
Coefficient
0.7
Target update period for MMD-based estimator
every 50 training steps
A.3 NOTIONS
For clarity and to ensure consistent notation, this paper adopts a set symbols throughout the train- ing and evaluation processes. The primary notations and their definitions are presented in Table 6
Table 6: Summary of the main notations.
Notation
Description
D
Offline dataset containing trajectories
{(s 02, 12)
A trajectory with states, actions, and rewards of length K. State, action, and reward at time step t. Return estimator output for state-action pair (st, at), represented N particles.
St. ai,
(st,
Parameters of the return estimator Zy.. Parameters of the target return estimator Z The i-th sampled particle from (st a:).
2r.())
Gr
Return-to-ge go at step d.
(("))
Optimistic value of (st, a:) computed by averaging top-n particles.
R
Aligned optimistic return from the trajectory starting point to step
t*
The peak step where Ri achieves its maximum.
Tttita
Subtrajectory sliced from step & fz. Set of unused timesteps during greedy slicing. Transformer-based policy parameterized by 0.
U
T
LMDD
MMD loss for training the return estimator. Training loss computed from the m-th selected subtrajectory. Total loss aggregated from all subtrajectory losses.
Lm
Liotal
7
Discount factor in Bellman updates.
N
Number of particles sampled by return estimator.
n
Number of top particles used for optimistic return computation.
A
Coefficient for weighting successive subtrajectory losses.
Retained history during evaluation.
V((s,)
Estimated optimistic value of current state time step t.
17

---

Published as a conference paper at ICLR 2026
B EXPERIMENT DETAILS
B.1 INTRODUCTION FOR BENCHMARKS
This study evaluates the proposed method on multiple standard benchmark tasks, covering continu- control, auction bidding, and instruction-driven navigation and manipulation. Screenshots of all tasks are provided in Figure 7
3
Agent
ADI
@ Goal
A/g
@ Boltuation
(a) D4RL.
(b) AuctionNet.
(c) BabyAL.
Figure 7: Screenshots of different benchmarks.
D4RL. D4RL (Datasets for Deep Data-Driven Reinforcement Learning) is widely used bench- mark suite for offline reinforcement learning. includes a variety of continuous control environ- ments, such MuloCo and Adroit control tasks (hopper, walker2d, halfcheetah, etc.), AntMaze and Maze2D navigation tasks, as well as Kitchen compositional manipulation tasks. A key feature of D4RL its provision of diverse datasets collected from expert, random, and geedium-performanee policies. enabling comprehensive evaluation of offline reinforcement learning algorithms under varying data quality conditions.
AuctionNet. AuctionNet simulates online advertising bidding, with data derived from large-scale auction logs. The bidding process is formulated as a reinforcement learning problem, where the agent's 's actions correspond to bid decisions, and the reward function is defined in terms of conver- sion rates and budget constraints. AuctionNet is commonly employed investigate policy learning under limited budgets, providing a testbed for assessing the adaptability and robustness of offline reinforcement learning algorithms real-world industrial applications.
BI AabyAl is an interactive wooldworld environment designed to evaluate agents in navigation and manipulation tasks driven by natural language instructions. The environment is partially observ- able and provides wide range of subtasks (e.g., picking up, placing, opening doors), which can be composed into more complex tasks. Its trajectory data includes state observations, actions, and lan- guage instructions, supporting research on the generalization ability of reinforeement learning and sequence modeling methods conditioned on instructions.
B.2 EXPERIMENTAL SETUP
In the D4RL environment, we employ the normalized score as the primary evaluation metric to en- able direct comparison with prior work. For AuctionNet, the evaluation focuses on cumulative con- version profit and budget compliance. We utilize the Alibaba simulation dataset from the NeurlPS 2024 competition (Xu et 2024). spanning 21 periods (7-27), each comprising over 500,000 im- pression opportunites and 48 auction steps, with each impression involving 48 bidding agents. The dataset contains more than 500 million records, including predicted conversion values, bid informa- tion, auction details, and impression outcomes. To ensure fair evaluation, data from periods 7-13 are used for training, while periods 14-20 are reserved for testing. For BabyAl, we adopt success rate the evaluation metric to assess the agent's ability to execute tasks under natural language instructions. The offline dataset is constructed using heuristic sampling, with 2000 trajectories gen- erated for each scenario to evaluate instruction execution and generalization capability. All reported results are averaged over three independent runs with standard deviations to mitigate the randomness introduced by different seeds.
18

---

Published as a conference paper at ICLR 2026
C ADDITIONAL RESULTS
C.I ADDITIONAL VISUALIZATION
The additional visualizations in Maze2D are shown in Figure S and Figure 9 In the maze-large environment, the original dataset comprises numerous trajectories spanning the entire map, with a substantial proportion failing to reach the goal. This leads to a mixture of high- and low-quality segments. Visualization of the return-to-go reveals that only trajectories approaching the goal exhibit high returns, whereas those farther away predominantly correspond to low-return regions. The MMD-based value estimator produces results that align with the overall trend of return-to-go but further distinguishes local value variations even in trajectories that do not directly reach the goal, thereby offering more fine-grained value signals. Leveraging these estimations, the PRGS sub-trajectory slicing mechanism prioritizes segments containing high-value regions, effectively discarding uninformative trajectories from exploration and substantially improving the quality of training data.
By contrast, the maze-umaze environment has a simpler structure, where trajectories are concen- trated in the U-shaped corridor, and most of them cover areas near the goal, yielding overall higher- quality data. Visualizations return-to-go and the estimator demonstrate a clearer return gradient in this environment, with estimated values strongly correlated with spatial positions. In such cases, although PRGS subtrajectory slicing introduces less dramatic changes than in maze-large, still ensures that training emphasizes critical high-value segments, thereby reinforcing the model's ploitation of key states while maintaining stability.
This comparison demonstrates that PRGS provides powerful filtering and segment selection in complex environments, while in simpler environments it functions as stable value-alignment mechanism.
C
E
E
(a) Trajectories.
(b) Rig visualization.
(c) Q estimation.
(d) Subtrajectories.
Figure 8: Visualization of trajectory data in maze-large.
(a) Trajectories.
(b) Rig visualization.
(c) Q estimation.
(d) Subtrajcctories.
Figure 9: Visualization of trajectory data in maze-umaze.
Within the hopper-medium environment, we perform t-SNE (van der Maaten & Hinton 2008 bedding and visualization of trajectory data to compare the distributions of return-to-g G and the MMD- Q estimates. As illustrated in Figure a),a), the visualization return-to-go reveals the diversity of trajectory quality in the dataset, where the point colors vary continuously from low
19

---

Published as a conference paper at ICLR 2026
values (purple) high values (yellow), indicating substantial mixture of high- and low-return samples within the same embedding space. In Figure 10jb), the Q-based estimates follow similar overall trend to return-to-go but delineate local structures more clearly: certain boundary regions estimated high-value, whereas these regions are less evident in the return-to-go visualization. This suggests that Q can recognize and emphasize potential high-quality state-action pairs in the modeling process, even when their actual returns in the raw data are obscured by noise or suboptimal trajectories.
Overall, these findings demonstrate that the MMD-based return estimator in the hopper-medium task not only captures the global value trends of trajectories but also identifies local value differences at finer granular- ity, thereby providing more discrimi- native criteria for subsequent subtra- jectory slicing and high-quality expe- rience selection.
(a) Rig visualization. Figure 10: Visualization of t-SNE embedding trajectory data in hopper-medium. Compared to the Rtg visualization, the Q estimation highlights more pronounced color differ- ences in the middle region.
(b) Q estimation.
C.2 ADDITIONAL ABLATION STUDY
To investigate the influence of X, which controls the weighting of successive subtrajectories, we conducted experiments in hopper-medium and maze2d-large, as shown in Figure [T The results indicate that the optimal value of À varies across different scenarios and algorithms. highlighting its sensitivity task characteristics. For example, hopper-medium (Figure I[)))..
( À = 0.7 yields the best performance for DT; while for BC, À = 0.3 leads to the most significant improvement. Never- theless, to better demonstrate the robustness of PRGS, we did not perform task-specific fine-tuning of X, but instead adopted unified setting throughout all experiments. It is also noteworthy that when À = 0, meaning that only the first sliced subtrajectory is considered, the method still achieves competitive performance, especially in maze2d-large. This observation further validates the key insight that explicit slicing of trajectories at the timestep level both meaningful and beneficial.
0,0,3
80
150
Performance
60
100
40
50
20
25
Algorithm
Algorithm
(a) Hopper-medium.
(b) Maze2d-large.
Figure 11: Influence of the À for PRGS.
We additionally conduct ablations focusing on the MMD- based return estimator. As shown in Figure 12. the full PRGS consistently achieves the best performance on both hopper tasks, whereas removing the optimistic MMD esti- mator (w/o Optimistic) replacing it with conventional Q-learning (w/ Q-loss) leads to substantial performance drops. Without the top-n particle selection, the return timates become overly conservative, weakening the ability to identify high-quality trajectory segments; replacing the
Normallzed Return
40
medium
hopper-medium replay
Gym Task
Figure 12: Ablation of the MMD
turn estimator.
20

---

Published as a conference paper at ICLR 2026
estimator with Q-learning further amplifies offline Q-value bias and significantly harms sequence modeling stability. These results demonstrate that the optimistic MMD return estimator an essen- tial component of PRGS, enabling reliable identification of high-value behavioral patterns and ultimately improving slicing quality and overall performance.
0.0
0.2
-0.
PRGS Optimistic PRGS Q-loss
PRGS
-0.6
-0.2
we Optimistic
-0.4
PRGS Q-less
0.0
0.2
0.6 Training Step
0.4
1.0 x10
0.0
0.2
0.4
0.6
0.8
1.0
Training Step
x10
(a) Hopper-mediun.
(b) Hopper-mediun-eeplyy
Figure 13: Bias (related to overestimation) comparison.
To evaluate whether the optimistic MMD-based return estimator leads to significant overestimation issues, we tested above three variants. The bias is calculated the difference between the estimated return and the true return:
Bias E[Q, Gr],
where Qe is the estimated return and Ge is the true return. As shown in Figure 13 the PRGS method does not exhibit noticeable overestimation and maintains a stable bias throughout training, compared to the other two variants, Empirically, a properly chosen target update period for the MMD-based estimator keeps the distributional estimates stable.
C.3 ADDITIONAL ANALYSIS
1.0
200
0.0
0,8
Prequency
Prequncy
3.50
-..6
300
1.00
200
100
0
0.0
0,0
15
20
First subttrjjectoyy length
First subtrajectory length
(a) Maze2d-large.
(b) Hammer-cloned.
1.0
600
1.0
Cumuuatiio Probabiitt
2.0
500
Prequency
350
Praquency
400
300
0.44
15
20
First subtrajectory length
First subtrajactory length
(c) Hopper-meddum-replayy
(d) Kitchen-complete.
Figure 14:
Distributions first subtrajectory lengths across different environments. The histograms (blue bars, left axis) show the frequency of the first subtrajectory lengths. while the orange curves (right axis) depict the smoothed empirical cumulative distribution functions (ECDFs).
Figure [4]pprrts the empirical distribution of the length of the first subtrajectory selected by PRGS across different benchmark environments. For each environment, about 2000 trajectories of length 20 are randomly sampled. The blue histograms denote the frequency of occurrence, while the or- ange curves represent the empirical cumulative distribution function (ECDF). The results show that in most environments the first subtrajectory length concentrated in relatively short range. For instance, in hanmer-cloned and kitchen kitchen-complets complete the proportion length-1 subtrajectories dom- inates, whereas maze2d-large and hopper-mediummreplly exhibit more dispersed distribution.
21

---

Published as a conference paper at ICLR 2026
Overall, short subtrajectories account for large proportion across tasks, indicating that the ini- tial slicing step tends to yield short segments. The variation across environments can be attributed to the inherent trajectory characteristics: in tasks with frequent high-reward signals or strong local cues, the first peak is often reached quickly, while in long-horizon navigation or locomotion tasks the return peaks tend appear later, leading longer subtrajectory slices.
1.0
1.0
1500
0.0
Frequenq
1000
-0.4
300
0.
121111111111111
7 101112131411161111
Number of penk-retum points per trajectory
Number of peakeretur points per trajectory
(a) Hopper-expert.
(b) Hopper-mediun-epper
0 protebbiliity
1200
1000
Frequeno
1000
Frequency
800
750
500
250
200
0.0
Number 101112331021611711111
Number
trajectory
(c) Hoper-medum--epllyy
(d) Hopper-medium.
1.0
4000
Traquendy
3000
2000
10000
Number of peakereturn points per trajectory
(c) Hopper-random.
Distribution of reak-return points per trajectory across different hopper datasets. The number of peak-return points varies significantly.
Figure 15:
Similarly, we also analyzed the number of peak- points. As shown in Figure 15 the number of peak-return points per trajectory varies significantly across different Hopper datasets. For the nopper-expert and nopper-medium-exper datasets, most trajectories have only a few peak-return points, typically around or 2, indicating that in these high-quality expert datasets, the return values tend to peak relatively fewer points during the trajectory. In contrast, the hupper-mediun-replle and hopper-medium datasets exhibit a greater number of peak-return points, with hopper medium- replay showing a more diverse distribution peak points across trajectories. Finally, the hopper- random dataset, almost all trajectories have peak-return points concentrated at 1 or 2, with most trajectories peaking early, reflecting the relatively low quality this dataset.
These results suggest that the quality and complexity of the dataset significantly affect the num- ber and distribution of peak-return points trajectories. High-quality datasets like Hopper-expert tend to generate peaks at fewer points, while lower-quality datasets like Hopper-random show more frequent early peaks.
C.4 ADDITIONAL RESULTS
As shown in Table 7 we also assess PRGS on narrow-coverage datasets such as random and expert. Since these datasets provide limited actionable coverage, with random datasets containing very few high-quality segments and expert datasets containing almost no low-quality segments to remove, the potential benefits of slicing naturally constrained. Nevertheless, PRGS still achieves small but consistent improvements under these restricted conditions, indicating its ability to adapt to different levels of dataset coverage.
22

---

Published as a conference paper at ICLR 2026
Table 7: Performance comparison on narrow-coverage Gym tasks.
Task
BC
BC-PRGS
DT
DT-PRGS
PDIT
PDIT-PRGS
halfcheetah--mng
-2.400.0
-2.4+0.0
-2.4+0.0
hopper- random
5.51226
0.3.10.5
6.7.1..0
8.111.3
7.6.10.6
7.6.10.8
walker2d-rundooe
11200
5.510.2
6.424..7
6.310.6
halfehoctall-cxpet
92. 02.55
93.2+0.1
919=0.1
93.6=0.3
90.1#1.1
93.2=10
hopper-expert
111.6 111.6+0.2
111.8+0.1
111.9+0.3
11.8-+0.4
111.640.22
111.+0044
walker2d expert
109.3.10.3
100.2:10.2
109.4:10.0
109.2.10.1
110.3:10.8
109.11111
Averaae
19.
14.
1012
To further examine the consistency between PRGS and actual task achievement, we additionally report a success metric for Maze2D, defined as the number of times the agent enters a target-near region within specified distance at any frame. This metric measures the success count within few episodes rather than success rate, and single episode may contain multiple successes. As shown Table S this success measure closely mirrors the return-based improvements across all methods, indicating that PRGS does not introduce ambiguity even when subtrajectories originate from failed or partially successful trajectories. Notably, PRGS selects subtrajectories based on optimistic return rather than instantaneous rewards; since return reflects the expected cumulative outcome toward the end the episode, it naturally aligns with this success-count metric, which is consistent with our empirical observations.
Table 8: Success count comparison within few episodes on Maze2D tasks.
BC-PRGS
DT-PRGS
PDIT-PRGS
maze2d-umaze
36.7=33.8
57.948.3
76.8+3.3
92.14:488
83.5+9.4
90.847.2
maze2d maze2d-medium
18 81:11.1
15.14332
81.655.3
188.4263
97.0#2.7
166 5116
maze2d-larg lurge
20.11:14.5
24.812.
64.216.4
221.3116.7
87.214.9
183.9110.1
Average
35.2
42.6
74.2
167.3
89.2
147.1
C.5 FAILURE MODE ANALYSIS: LIMITED TRAJECTORY COVERAGE
In addition to the main experiments, we analyze a key failure mode of PRGS that arises when the dataset lacks sufficient actionable structure. As shown in Table 7 evaluate PRGS narrow- coverage datasets such as random and expert settings, particularly in expert datasets. These datasets expose PRGS to two extreme conditions:
Only a few isolated high-return segments exist, and most trajectories contain no meaningful hogghorizonn structure.
Random datasets:
Expert datasets: Trajectories are already near-optimal, and contain almost no low-quality segments.
Under both conditions, the potential benefit of slicing is naturally constrained. Since PRGS relies on identifying high-value transitions and separating them from low-quality ones, limited structure reduces the number of actionable peak-return points and, consequently, the effectiveness of the slicing procedure.
In these narrow-cove coverae regimes, PRGS behaves similarly baseline methods: the lack of usable structure prevents the construction meaningful composite trajectories, and the method is unable perform effective stitching.
D COMPUTATIONAL OVERHEAD
As shown in Table the computational overhead of PRGS comes primarily from the MMD-based return estimator, whose matrix operations introduce moderate additional memory usage during train- ing. The greedy slicing procedure also slightly reduces training speed, although it is applied only during training and has effect inference. In our profiling, PRGS lowers training throughput by about 16% and increases memory usage by roughly 89%, while maintaining nearly identical infer- ence cost. These overheads are acceptable given the consistent performance improvement brought by PRGS, and they mainly reflect the Eier-graines value estimation and data selection that underpin its effectiveness.
23

---

Published as a conference paper at ICLR 2026
Table 9: Computational overhead of PRGS during training and evaluation.
Method
Training Throughput (1)
Memory (4)
Relative Gain (t)
DT
1x
1 x
1x
PRGS (training)
0.84
1.89
1.15x
PRGS (evaluation)
0.99x
1.08 x
LLM USAGE
Large Language Models (LLMs) were used only as an assistive tool for grammar and language polishing of the manuscript. They did not contribute to research ideation, methodology design, experimental implementation, analysis, or conclusions.
24


Reviews:
[{"review_id": "2zPM0PaLTa", "reviewer": "Reviewer_oYSY", "summary": "The paper presents Peak-Return Greedy Slicing (PRGS), a framework for improving the learning process of Transformer-based architectures in offline deep reinforcement learning. The main idea is to identify high-quality sub-trajectories and use them for higher quality training of the model. The sub-trajectories are chosen by identifying points in the trajectory where the overall reward (obtained+predicted) begins to decline. The authors use a mechanism to \"reconcile\" trajectory history with the sub-trajectories they use. Extensive evaluation on multiple datasets and environment types shows a generally consistent improvement.", "questions": "I invite the authors to respond to the points I raised in the \"weaknesses\" section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The timestep-level sub-trajectory slicing approach is novel, simple and elegant. The analogy to a \"human\" way of learning is clear.\n2) The approach is relatively model agnostic, and can be used in multiple offline methods.\n3) The evaluation is comprehensive, consisting of multiple environments and algorithms.", "weaknesses": "1) There is no discussion of the computational complexity of the approach. The slicing and MMD estimator are likely to require non-trivial computational resources.\n2) The subject of training trajectory quality is likely to be a key element in the success or failure of the proposed approach. The lack of discussion and analysis of this point is problematic.\n3) There is a lack of analysis regarding to how often the proposed approach actually takes place: in what percentage of trajectories does slicing take place? what is the average length of a sliced trajectory? what is the impact of history truncation?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "9zlw6mdoCo", "reviewer": "Reviewer_5VUZ", "summary": "The paper proposes a framework to enhance the performance of transformer-based offline RL methods. Peak-Return Greedy Slicing, helps enable discovery of high quality subtrajectories, and in turn do BC on those subtrajectories. The method first trains a Q function with MMD, which approximates the distribution of Q values. Then, a trajectory is partitioned into subtrajectories by recursively determining peak Q values, and then do BC on each subtrajectory. During inference time, the history is truncated whenever the value increases.", "questions": "- why is the sum of return not discounted in the definition for “aligned optimistic return” at the bottom of page 4? The Q values are learned with discount, so doesn’t this favor returns toward the start of the trajectory?\n- in section 5.3, how do you filter for the top 10% or 20% of the trajectories?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- the paper is very well written, and does a great job at describing the method and walking through the experiment results. The plots are well made to help with understanding the paper.\n- The experiment section provides extensive number to validate PRGS, and shows the method works well across diverse benchmarks. There are also sufficient baseline to compare the method against.\n- I appreciate the visualization in section 5.4, which help with intuition of the method", "weaknesses": "- It would be great to offer some discussion on the failure modes of this method. For example, consider a subtrajectory where there are multiple part where the values go down, but eventually the values rise to the “peak” of the episode. When we do BC on this trajectory, don’t we then inherently learn to “clone” suboptimal behavior? How does PRGS deal with this?\n- it is unclear why we need optimistic Q values. Traditionally, offline RL Q function needs to be conservative to prevent overestimation. Figure 4 does show empirically that optimistic Q values seem better. But since this seemingly goes against traditional wisdom (though I can think of why it’s different for transformer-based offline RL), it would be nice to include some intuition and explanation in the paper.\n- it is unclear to me why in Section 3.3, the authors chose to remove history when the value increases. During training, doesn’t the value of mostly increase for each subtrajectory? Then doesn’t this cause a train-test mismatch, when at training time history is kept when values increase and at test time history is removed when values decrease?\n- nitpick: section 2.2 could have highlighted more the difference between transformer-based offline RL and traditional offline RL", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "BfuLjNqS1c", "reviewer": "Reviewer_EfRn", "summary": "The paper introduces Peak-Return Greedy Slicing (PRGS) to enhance Transformer-based Offline Reinforcement Learning. PRGS uses an MMD-based return estimator and greedy slicing to extract high-quality subtrajectories instead of relying only on complete ones. This method significantly improves performance by better exploiting and recombining valuable data segments.", "questions": "See weaknesses,", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) Well motivated. The authors propose to utilize high quality subtrajectories to improve performance of transformer based RL.\n2) Results seem good. The proposed method shows strong improvement on reward.", "weaknesses": "1) The core strength of the proposed method appears to be its ability to capture and utilize peak intermediate rewards within subtrajectories, leading to the reported strong average reward improvements. However, the use of subtrajectories that might originate from ultimately failed full trajectories introduces a potential ambiguity. Since the primary metrics presented in the paper are based on reward, it would significantly enhance the clarity and robustness of the results if the authors could provide the success rate for the respective tasks.\n\n2) As depicted in Figure 3, the training procedure involves masking low quality tokens, leading the Transformer to learn intermediate high quality tokens conditioned on low quality segments. A potential concern is that this mechanism might inadvertently encourage reward hacking: the policy may learn to exploit the high intermediate reward signal without consistently leading to final task success. Specifically, the model might optimize for peak reward segments from trajectories that ultimately terminated in failure. How about to implement a filtering mechanism to exclude subtrajectories (state-action pairs in the env) that can not reach the final success?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "z7fpgOgud6", "reviewer": "Reviewer_qNL9", "summary": "The manuscript proposes Peak-Return Greedy Slicing (PRGS), a framework designed to enhance Transformer-based offline reinforcement learning (RL). The method introduces a fine-grained, timestep-level trajectory segmentation mechanism that selects high-quality subtrajectories instead of relying solely on full trajectories. PRGS consists of three components: (1) an MMD-based return estimator that models the distribution of potential future returns for each state-action pair and yields optimistic estimates (Sec. 3.1, Eq. (2–4)); (2) a greedy slicing algorithm that recursively partitions trajectories based on peak optimistic returns (Sec. 3.2, Eq. (5)); and (3) an adaptive history truncation mechanism for evaluation, ensuring alignment between training and inference (Sec. 3.3, Eq. (6)). Experiments on diverse benchmarks (D4RL, BabyAI, AuctionNet) demonstrate consistent performance improvements, averaging 15.8 % gains over Transformer-based baselines (Sec. 5.1; Table 1).", "questions": "- How does PRGS perform on datasets with narrow coverage, such as random or expert datasets?\n- Why are the comparisons with QDT, EDT, and CGDT conducted only on the medium-replay and medium tasks, rather than across all datasets?\n- How biased is the return distribution estimation? Could it lead to overestimation? Would optimistic estimation make the overestimation even more severe?\n- How does PRGS perform in terms of time and GPU memory consumption? Could you provide a detailed comparison with other algorithms?\n- What would the performance be if we do conventional Q estimation rather than estimating the return distribution?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The writing and presentation of this paper are intuitive and clear. The paper effectively explains the motivation behind the method’s design. I can clearly understand the three components of the approach — the evaluation of the return distribution, the trajectory slicing mechanism, and the adaptive history truncation during deployment — as well as the rationale behind each of them.\n- The experiments in the paper are thorough and align well with the authors’ claims. Applying PRGS to Transformer-based offline methods such as DT yields noticeable performance improvements, demonstrating the effectiveness of the proposed approach. Some ablation studies, such as examining the impact of the number of particles used in return distribution estimation, also correspond well with the authors’ design intentions.", "weaknesses": "- Although PRGS improves the performance of Transformer-based offline algorithms, its overall performance is not particularly outstanding. The compared RL-based baselines include only CQL, IQL, BEAR, and TD3+BC, which are not among the most recent RL-based methods. To my knowledge, the state-of-the-art performance in offline RL is significantly higher than what is reported in the paper.\n- Intuitively, the effectiveness of the PRGS method is likely to depend on the diversity of the dataset. If the data coverage is narrow—for example, consisting entirely of expert demonstrations—the number of peak-return points should be relatively small. The paper also lacks experiments on random and expert tasks, and the authors do not show how different types of datasets affect the number of peak-return points.\n- Because dataset coverage is often incomplete, value estimation in the offline setting typically suffers from significant bias and even overestimation. Although the paper estimates the return distribution rather than a single value, this issue is still likely to occur. The authors, however, do not appear to provide any analysis or mitigation of this problem.\n- PRGS is somewhat more complex than standard Transformer-based methods, yet the authors do not provide any experimental comparison of its computational cost in terms of runtime or GPU memory consumption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)