Rebalancing Return Coverage for Conditional Sequence Modeling in Offline Reinforcement Learning
Authors: Wensong Bai
Chufan Chen
Yichao Fu
Qihang Xu
Chao Zhang
Hui Qian
Venue: NeurIPS 2025 poster ()
Keywords: Offline Reinforcement Learning
Conditional Sequence Modeling
Return Coverage Rebalancing
Primary area: reinforcement_learning
TL;DR: We reveal how return-coverage affects the performance of conditional sequence modeling policies in offline RL and propose an algorithm achieving new state-of-the-art results on D4RL.

Abstract: Recent advancements in offline reinforcement learning (RL) have underscored the capabilities of conditional sequence modeling (CSM), a paradigm that models the action distribution conditioned on both historical trajectories and target returns associated with each state. However, due to the imbalanced return distribution caused by suboptimal datasets, CSM is grappling with a serious distributional shift problem when conditioning on high returns. While recent approaches attempt to empirically tackle this challenge through return rebalancing techniques such as weighted sampling and value-regularized supervision, the relationship between return rebalancing and the performance of CSM methods is not well understood. In this paper, we reveal that both expert-level and full-spectrum return-coverage critically influence the performance and sample efficiency of CSM policies. Building on this finding, we devise a simple yet effective return-coverage rebalancing mechanism that can be seamlessly integrated into common CSM frameworks, including the most widely used one, Decision Transformer (DT). The resulting CSM algorithm, referred to as Return-rebalanced Value-regularized Decision Transformer (RVDT), integrates both implicit and explicit return-coverage rebalancing mechanisms, and achieves state-of-the-art performance in the D4RL experiments.

Rebalancing Return Coverage for Conditional Sequence Modeling in Offline Reinforcement Learning
Wensong Bai', Chufan Chen', Yichao Fu', Qihang Xu', Chao Zhang 2 Hui Qian'
CCollege of Computer Science and Technology, Zhejiang University, Hangzhou, China Advanced Technology Institute, Zhejiang University, Hangzhou, China {wensongb, chufan. chefn, chen, fuyichao, 22351233, zczju, qi anhui) 1 @zju.edu.cn
Abstract
Recent advancements in offline reinforcement learning (RL) have underscored the capabilities of conditional sequence modeling (CSM), a paradigm that models the action distribution conditioned on both historical trajectories and target returns associated with each state. However, due to the imbalanced return distribution caused by suboptimal datasets, CSM is grappling with serious distributional shift problem when conditioning on high returns. While recent approaches attempt to empirically tackle this challenge through return rebalancing techniques such weighted sampling and value-regularized supervision, the relationship between re- rebalancing and the performance of CSM methods is not well understood. In this paper, we reveal that both expert-level and full-spectrum returncoverrge crit- ically influence the performance and sample efficiency of CSM policies. Building rebalancing mech-
on this finding, we devise a simple yet effective return-coverte anism that can be seamlessly integrated into common CSM frameworks, including the most widely used one, Decision Transformer (DT). The resulting CSM algo- rithm, referred as Return-rebalanced Value-regularized Decision Transformer (RVDT), integrates both implicit and explicit return coverage rebalancing mecha- nisms, and achieves state-of-the-an performance the D4RL experiments.
Introduction
Offline RL aims learn effective policies solely from previously collected datasets, without any further environment interaction [40 (17, 25, 30J. Recent advancements in offline RL primarily fo- cus on the use of conditional sequence modeling (CSM). where the core idea to learn the return- conditioned distribution of actions at each state given historical information [6]]] 25,, 65,, 66, 26]. One of the most representative methods in CSM the decision transformer (DT) [[]] wherein historical trajectories composed of multiple return-state-actinn based model. During inference time, the model predicts sequence of actions conditioned on a target return, effectively reframing the offline RL problem as a supervised learning task. By enabling the policy make decisions directly from temporally extended trajectories, this paradigm circumvents the need for recursive value estimation, thereby avoiding both the challenges of Bellman complete- methods [III, 68).
triplets are exploited to train a Transformer-
ness and the bootstrapping errors associated with dynamic programming-bant
However, CSM-based algorithms encounter substantial difficulties when learning from suboptimal datasets, particularly due to the imbalance between near-optimal and low-return trajectories within datasets [5 42]. This uneven distribution of trajectory quality in the training dataset may induce
*Corresponding author: zczju@zju. edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

significant distributional shift problem at inference time when conditioning on high returns, critically
limiting the achievable performance of CSM methods [T TQ]. Recent advances have acknowledged such difficulties and proposed various mitigation strategies by rebalancing return distributions [A8,
23, 92, 57, 25]. For instance, DiffStitch [#2] employs diffusion models to synthesize high-return
trajectories by stitching together optimal sub-segments drawn from multiple suboptimal trajectories.
Critic-Guided Decision Transformer (CGDT) [ and Q-value Regularized Transformer (QT) [25]
integrate value-based guidance into the DT framework, implicitly assigning greater importance to near-optimal actions, thereby steering policies toward higher-return regions. These methods achieve enhanced performance through explicit implicit rebalancing of the return distribution, and their efficacy has been empirically validated by achieving state-of-the-ar (SOTA) results [ 57,, 25]
Nonetheless, none of these prior works provides a theoretical analysis elucidating the fundamental principles underpinning their empirical successes.
In this paper, we uncover the relationship between the return distribution of the training dataset and the performance of CSM policies at inference. Specifically, two aspects of return-coverage are relevant to the performance: (i) expert-level return-coverage, corresponding to the coverage of near-optimal and expert-level returns; (ii) full-spectrum returncoveraag, which corresponds to the retum-coverage over the entire return space By analyzing the coverage of these two returns within the offline datasets, we derive theoretical bounds quantifying the performance gap between the learned CSM policy (given fixed inference- efference-time conditioning functions) and the optimal policy T of the underlying MDP. Our analysis reveals that the performance gap is jointly influenced by both expert-level and full-spectrum retum-coverage, providing principled explanation for why return distribution rebalancing enhances policy performance and mitigates distributional shift at inference.
Building upon this finding, we propose a simple yet effective return-coverage rebalancing mecha- nism, which formulates rebalancing KL regularization between the agent's 's policy and an expert policy extracted from near-optimal trajectories in the dataset. The proposed mechanism explicitly interpretable and seamlessly compatible with most existing CSM algorithms as plug-in module. Leveraging this mechanism, we further design a practical CSM algorithm named Value-regularized Decision Transformer (RVDT). which achieves consistent performance improve- ments over state-of-the-an methods. Our contributions are summarized as follows:
1. We theoretically show how the performance of CSM-based policies is influenced by the coverage of expert-level and full-spectrum conditional returns the offline training dataset. Specifically, we show that the gap between the expected return of the learned CSM policy and the optimal return is inversely proportional to the coverage of these two returns, and we derive an upper bound to explicitly characterize the associated sample complexity.
2.
We propose an explicit return-coverage rebalancing mechanism as a plug-in module for enhanc- ing existing CSM algorithms. In practice, an independent network is trained via imitation learn- ing on expert-level trajectories to extract expert policies, whose KL divergence with the agent's s policies used regularizer. This regularizer has been shown perform weighted sampling explicitly over trajectories when the policy is parameterized by a Gaussian distribution.
3. We design RVDT, which utilizes both explicit and implicit rebalancing mechanisms, effectively addressing the difficulties posed by limited expert data through the utilization of an expert KL divergence regularizer and a value-regularized guidance. This combination collectively mitigates inference-time distributional shift, enabling superior performance in challenging scenarios such as sparse-reward reward environments and low-quality datasets.
The performance of RVDT is evaluated on the D4RL. benchmarks [16]. Experimental results demon- strate significant performance gains over representative baselines, establishing RVDT as a strong candidate for advancing the state-of-the-art (SOTA) in offline RL.
2 Related Work
Learning from imbalanced datasets, comprising predominantly low-return trajectories and relatively few high-return trajectories, has long been a critical challenge in offline RL [23, 21, 14, ICJ. Broadly,
'The full- alll-spetrum return- coverage is closely related to the runtime returns, i.e., the realistic returns during policy execution, as the actual distribution of runtime conditional returns during execution is typically uncertain.
ii

---

prevailing methods tackling this challenge can be categorized into two types: explicit rebalancing and implicit rebalancing. Explicit methods encompass techniques such as dataset selection [I2
48, [2], weighted sampling [44, 211, 23), and data synthesis |62
34
E2I, while implicit methods
generally leverage value estimation as additional guidance [38
D5 17
25).
Sub-dataset Selection. Several methods explicitly select informative subsets from imbalanced datasets, e.g., sample uncertainty [I2, ]8] or frequency of being forgotten [bai] is utilized to prune the datasets. Other studies employ submodular functions that capture the notions of diversity, repre- sentation, and coverage to eliminate redundancy [52,, 60 59 32]. A more straightforward approach, top-% BC, directly selects and fits only the highest- trajectories within the dataset [69 [T].
Weighted Sampling. Another line of research adopts weighted sampling techniques to rebalance trajectory distributions [24, II, 22 _
Hong et al. [2T] leverage a Boltzmann distribution based on trajectory returns advantages reweight trajectories. Similarly, density-ratio weighting [2X]
utilizes importance sampling [49] to emulate sampling from other supplementary datasets.
Data Augmentation and Generation. Other methods address trajectory imbalance by augmenting or synthesizing data [[56, 35 5-1, 62, 12 For instance, QDT [62] relabels return tokens using
a learned conservative Q-function, subsequently integrating these augmented trajectories into the training datasets. DiffStitch [#2]] leverages diffusion models to generate high-return trajectories by
stitching optimal sub-segments extracted from multiple suboptimal trajectories.
Value-regularized Supervised Learning. Recent works such as BEAR [X CDC (USS], O-RL. [Z]. TD3+BC [0 TD3+RKI. [9], CGDT [5.2], and QT [25] integrate value-based guidance into supervised learning frameworks to prioritize high-return actions. As a result of leveraging transition- level value estimates to prioritize near-optimal actions without explicitly modifying the trajectory
return distribution, these algorithms are considered as implicit return rebalancing methods.
3 Preliminary
3.1 Offline Reinforcement Learning
We model an agent-enviromment interaction as a finite horizon Markov Decision Process (MDP) (S,A, T, Mo. H), where S and A denote the (continuous) state and action spaces. R S - R is the reward function, T defines the transition dynamics, i.e., $1+1 ~ T az). Ag is the initial state distribution, i.e., 81 ~ mo. and H is the time horizon. In the offline RL setup, algorithms given a static dataset 'D consisting of trajectories = {81, (1,771," H an, TH), collected by a behavior policy тT. Let g(Tz) >d Ti (abbreviated as gr) denote the return-to-go (RTG) of trajectory at time step I, and J (T) = Exar [g())] is the expected return of a policy x.
3.2 Return Conditioned Sequence Modeling
Motivated by the wide-ranging applications and remarkable success of transformer models in var- ious domains [55L & I., 27 29 36), CSM-based methods formulate policy learning as a return- conditioned supervised learning problem. Specifically, during training, the policy и is optimized by minimizing the empirical negative ligeliieldoood (NLL) loss over dataset DD of trajectories:
[(T) = -EEoDD log
g( 7tK
(1)
where -.K = [gt-KK K St-K, K(-K...., 91-1 5t-1, 1-1} denotes the historical input with context
length K. During inference, the learned policy π is combined with conditioning function f(s) to form the inference-time policy π as T (a|s) ;= x (a]s, f (s), F*), where s denotes the collection of and historical triplets. Algorithms such as DT and its variants, e.g., QT [2S] and CGDT [57]] can be viewed as specific instantiations within this framework.
The NLL loss can be interpreted as approximating the conditional distribution Pr,, (als. where Pro denotes the distribution over states, actions, and returns induced by the behavior policy wy- In the regime of infinite data, it has been shown that the CSM objective effectively reweights the behavior policy based on the distribution of future returns. Under this setting, the optimal return- conditioned policy of
CSM for given conditioning function f can be expressed as: 7CSM (als) =
iii

---

(a(a)s)) Pro (M()))s.)) Pap ([+))s With infinite data and a fully expressive policy class, optimizing (I) guarantees that the learned policy rrfM is provably near-optimal [K]].
Theorem 1. Consider a finite-horizon MDP with horizon H, behavior policy Tg. and condi- tioning function f. Suppose the following conditions hold: (i) The return is sufficiently covered, i.e., (g = f(s1)(s)) 2 af for all initial states $1. (ii) The environment is nearly deterministic in the sense that for some functions TT and R, i.e., P(r * R(s,a) a) or + T(s. a))s, a) S e, Vs, (iii) The conditioning function f is consistent such that f(s) = ((s') + r, vs. Then the policy that optimizes the CSM objective (Eq. W is guaranteed to be near-optimal: Ez, f(1) (TT S (1/ag + H
Theorem DI establishes theoretical guarantees and success conditions for CSM methods. Brandfon- CSM and x², where I corresponds
brener et al. [6] further analyze the performance gap between क f to the optimal conditioning function. However, at inference time, the runtime conditional returns, computed as the prior returns minus sampled rewards r, cannot consistently remain optimal due to the inherent randomness of r [S) [3]. This discrepancy between theory and practice substantially constrains its practical utility for enhancing the robustness and generalization of CSM algorithms.
4 Methodology
This paper focuses on investigating how return distribution of the training dataset impacts the per- formance of CSM policies through the lens of return-coverage. We theoretically reveal that the performance gap between the MDP's optimal policy л and the learned CSM policy is jointly de- termined by both expert-level and full-spectrum return-coverage. Likewise, we establish that the sample complexity is similarly governed by these two aspects of coverage. The theoretical insights elucidate why previous approaches, utilizing explicit or implicit rehalancing of return distributions, effectively improve performance. Following this notion, we propose a simple yet effective return- coverage rebalancing mechanism, which is explicitly interpretable, and seamlessly compatible plug-in module with most existing CSM algorithms. To demonstrate the effectiveness of this mech- anism, we integrate it into QT [25], the current SOTA, resulting in a new approach termed RVDT.
4.1 Analysis on Return-Coverage Rehalancing for Conditional Sequence Modeling
Due to the imbalance between expert-level and low-rcturn trajectories in offline datasets, the study of Brandfonbrener et al. [6], which focuses on the optimal conditioning function f", becomes imprac- tical for real-world scenarios. Our analysis extends this perspective by considering a more general function class that encompasses both potentially non-optimal runtime conditioning functions and optimal (or near-optimal) conditioning functions, denoted as f and f*, respectively. The runtime conditional return function f corresponds to arbitrary possible returns that may be encountered dur- ing policy execution. Let 9 denote the collection of all possible returns collected by π E II, then f:S G. The target conditional function f that we aim to find is the RTG under the optimal policy x., i.c., f(s) max, E, [g])]. We define the expert-level return-coverage in the offline dataset generated by behavior policy Tg as (g = f (s1) where the f* corresponds to the optimal policy x" of the underlying MDP. Correspondingly, the full-spectrum return-coverage is defined as PT (g f((1)] s1). We first analyze the performance of CSM policies under an ideal- ized setting with infinite data and fully expressive policy class, establishing an upper bound on the suboptimality of CSM policies and quantifying its dependence on both aspects of return-coverage
Theorem 2 (Performance gap with respect to return-coverage, Consider a finite-horizon MDP with horizon H, behavior policy πB. a runtime conditioning function f. and the optimal conditioning function for n" is fr. Assume the following assumptions hold: (i) Return-coverage: Pag (g [(s))ss) of and Pxg (g= ("())sss)))) Z ož for all initial states 81. (ii) Near determinism: P(r + R(s,a) a) s' + T(ss, a) s, a) < e at all (s,a) for some T and R. (iii) Consistency of f: f(s) f(s') + r for all s. Then the following upper bound holds:
1 J(T") - J(n < aj
+3)H22
(2)
aj 0z
where C' E (0. 1) is constant. The proof is provided in Appendix B
iv

---

The result formally characterizes how the effectiveness of policies trained via CSM is constrained by both expert-level and full-spectrum returncoverage. This implies that algorithms are capable of modulating o? and of through rehalancing the return distribution, thereby facilitating the conver- gence of C'SM toward the optimal policy πτ. Notably, in the offline RL setting, the relationship of + az 1 holds when f * f" In other words, strategies such as sub-dataset selection or weighted sampling that attempt to increase one type of retur-coverage without altering the distribu- tion's ss support tend to reduce the other.
The subsequent analysis examines the sample efficiency of CSM methods under common assump- tions as in [6], with finite data and a restricted policy class. Letting N denote the size of the dataset, Theorem B demonstrates that the sample complexity exhibits structural dependence on the return- coverage analogous to that Theorem Note that for clarity, we consider the case of single 8 (7 Ø)) here, which can be readily extended to the non-empty - scenario.
Theorem 3 (Sample complexity). To get finite data guarantees, add to the above assumptions in Theorem Q that (i) bounded occupancy mismatch: P Pappu()) S Cy PPg (s) for all s; (ii) the policy class II is finite; (iii) [log (a) s, 9) log ι (a')s',g') S e for any (a.8. 8. 9. a', ,'',g) and all E II; (iv) the approximation error of II is bounded by Capprox , L.e, min L(TT) S Capprox.
Define the expected loss as L(#t = Ea-Pag Pyr~Pag(1e (-le) [KL (Pa (118,) ) || ( 9))] Then for any estimated CSM policy f that conditions f at inference time, with probability at least 1 1- 6,
[
C a
log (T//88 N
1/4
C1 as
E+C C
C
H² m
J( ( so
The proof. provided in Appendix 8 utilizes uniform convergence results from supervised learning theory [SD].
Theorem B quantifies the sample complexity of CSM methods in terms of expert-level and full- spectrum return- -ooverae Together with Theorem E1 our analysis unveils the fundamental principles underpinning the empirical successes of current return-coverage rebalancing strategies. That is, rebalancing a and a reduces the gap between the learned CSM policies and x².
4.2 Plug-in Return-Coverage Rebalancing Mechanism for CSM
Building upon the theoretical insights established in preceding theorems, we propose an explicit re- balancing mechanism based on KL divergence regularization with respect to a learned expert policy T. Specifically, a subset of near-optimal trajectories is filtered from the offline dataset, denoted De C DD from which n" is extracted via imitation learning. Given CSM policy Tо parame- terized by DT, our proposed rebalancing mechanism operates as plug-in module appended to the original ÑLL loss in Eq. MI We refer to the resulting CSM algorithm as Return- rebalanced Decision Transformer (RDT), which is optimized with the following loss:
H Σ [T,[To 9(T,), (l (4)
(R) = EE
Σ 1=1
log To(al s(T,), +aE +
The first term in Eq. (4) aligns the model with the empirical action distribution of тg, while the regularization term guides the policy towards expert-level actions. Notably, the rebalancing mod- ule adopts KL divergence instead of MSE, designed to maintain compatibility with broader policy classes in modern CSM methods, such as diffusion models [6% A7]].
To illustrate how the expert KL regularizer explicitly rebalances the return-coverage, we demon- strate that under a common stochastic policy parameterization, the regularization process can be equivalently interpreted as trajectory-level reweighting of expert-level samples in the offline dataset. Proposition 1 (KL regularization as weighted sampling strategy). Assume the policy na is param- eterized by factorized Gaussian distribution with a fixed standard deviation. Then optimizing the RDT objective in (H) is equivalent to optimizing the following weighted NLJ. loss (proof in A.4):
arg min Exor (0) arg min E. T4
(1 a -1([r E D..) (E log No ra(ai|sa, (a)))))s.(TT)) (()))
(5)
Proposition M shows that the KL regularization effectively performs a weighted resampling of ex- pert trajectories, practically increasing the expert-level return-coverage within the offline dataset. Though an optimal rebalancing ratio exists, it is generally infeasible to compute explicitly. In prac- tice, our rebalancing mechanism, akin to other rebalancing strategies [177, H. 211], amplifies the gra- dient contributions of expert-level behaviors during training.

---

4.3 Return-rebalanced Value-regularized Decision Transformer
We integrate our proposed rebalancing mechanism into the current SOTA method. QT, resulting in novel algorithm that incorporates both explicit and implicit return-coverage rebalancing mecha- nisms, referred as RVDT. By implicitly rebalancing the relurncoverage with Q-value regularizer to DT, the loss function of QT' is defined as follows.
(or = B,-D [E == - log To (1 (
- -
(s1, (6)
However, accurate Q-value estimation can be prohibitive in complex environments, limiting the effectiveness of QT"s value-based rebalancing mechanism. In contrast, RVDT uniquely combines explicit and implicit rebalancing mechanisms, i.e., an expert KL divergence regularizer and an value- regularized guidance, enhancing its adaptability to challenging offline datasets. The resulting learn- ing objective for RVDT augments policy regularization with value-bace based guidance:
LRVITE (0) = EE.DD E iill log ma(a s., 9(7),, o. nE Brr ma 1Q
(s:: as)]
(7)
a., D. E. KL [ (T.)), (1==))l(""((() (+]s))]
We adopt the same parameterization strategy as QT and employ offline DP to iteratively learn the value function QTe The coefficient 17 follows commonly adopted normalization scheme [[]]], i.c.
77 = Er-wE ", (=,,e)]) where 1j is a tunable hyperparameter. The training and inference details, along with the algorithm pseudocode, are provided in Appendix H.
Compared to QT, RVDT not only inherits its strengths but also benefits from directly aligning its pol- icy distribution with expert-level trajectories, thereby significantly enhancing performance in sparse- reward tasks where accurate Q-value estimation is infeasible. While empirical evidence (Section 5 confirms the superiority of RVDT, we provide a theoretical perspective to substantiate its superior- ity over related methods under idealized conditions (i.e., infinite data and neglecting approximation errors). The formal statement and proof of Proposition Q are provided in Appendix A.S.
Proposition 2 (Informal). Let n.. Tyo and Tे. denote the optimal policies learned from the objec- tives of RVDT, RDY, and QT, respectively. Given the behavior policy n3 there exists an expert-level subset De S DD and an appropriate coefficient a 2 O such that the following inequalities hold for all stutes 8 ~ drgi VVV" (8) 2 E. V (s) Eg V (8) and E, V (s) 2 E, V TT (s) 2 E2 v (5).
5 Experiment
We conduct a comprehensive empirical evaluation of our proposed approach on the D4RL bench- marks [[ RVDT is compared against broad range of baselines, mainly composed of DP DP-based
methods [37] 19, 20J, value-regularized SL [38 I1, I77, and CSM approaches [1II] 51 62 28, 577
63, 25). To further assess the robustness of RVDT, we design a set of more challenging datasets to compare RVDT with the current SOTA method (QT) [25]., focusing on three key desired properties: (i) effective utilization of limited expert data, (ii) learning ability under low-data regimes [H66), and (iii) trajectory stitching ability. Finally, we perform carefully designed ablation studies to isolate and examine the contributions of key components in RVDT, including the KL regularization-base return-coverage rebalancing mechanism, the Q-value guidance, and the stochastic policy modeling.
Benchmarks and Baselines. We evaluate method across five diverse domains from the D4RL benchmarks [IG]]: Gym, Adroit, Kitchen, Maze2D, and AntMaze, encompassing tasks from continu- control multi-task manipulation and -orizon planning. Our approach is compared against a broad set of baselines, including CSM-based methods (DT [[I], StAR [5]]] QDT [62),, GDT [28], (25] ), Da-based methods (IQL [37]. BCQ [IT]] CQL [OO), value-regularized SL
CGDT [5 QT
methods (BEAR [X O-RL [ TD3+BC [[Z]), and diffusion-basse methods (Diffuser [G]]]. De-
cision Diffuser [2]. Diffusion-OO QL [58]). We further include model-based approach (MoRel [{3]]
for broader comparison. QT is reproduced using its official implementation, while other basclines follow their original publications or results reported in QT's paper, depending on availability. Addi- tional benchmark and implementation details are provided in Appendix C. and DI
5.1 Main Results on D4RL Benchmarks
The experimental results comparing RVDT against baselines on the D4RL benchmarks are presented in Table RVDT consistently outperforms most baselines across various domains, establishing new SOTA performance on the majority of D4RL tasks. Tesk- comparisons provided below.
vi

---

Table 1: Performance comparison across D4RL tasks. For each task, RVDT reports the mean and
standard error of normalized scores [16] over 30 random rollouts and averaged over 5 random seeds.
Gym Tasks
CQL
IQL
BCQ
TD3+BC
MoRel
BC
DD
DT
SLAR
GDT
CGDT
QT
RYDI
hallcheetah me
91.6
86.7
69.6
90.7
53.3
55.2
90.6
86.8
93.7
93.2
93.6
93.2
94.4 1 0.1
hopper-m-e
105.4
91.5
109.1
98.0
108.7
52.5
111.8
107.6
111.1
111.1
107.6
113.0
113.1
+
walker2d-m-e
108.8
109.6
67.3
110.1
95.6
107.5
108.8
108.1
109.0
107.7
109.3
112.0
112.7
halfchectal-m
49.2
47.4
41.5
48.4
42.1
42.6
49.1
42.6
42.9
42.9
43.0
51.0
51.9
- 0.3
hopper-m
69.4
66.3
65.1
59.3
95.4
52.9
79.3
67.6
77.1
99.6
100,2
+ 0.1
walker2d-m
83.0
78,3
52.0
83.7
77.8
75.3
82.5
74.0
73.8
76.5
79.1
87.2
90.2
+ 0.1
hallcheetah m-r
34,8
36.6
36.6
40.5
48.8
2.0
hopper-mes
95.0
94.7
31.1
60,9
93.6
18.1
100.0
82.7
29.2
85.3
93.4
102.1
103.2
1.9
walker2d-m-r
77.2
73.9
13.7
81.8
49.8
32.3
75.0
79.4
39.8
77.5
78.1
97.8
99.3
0.8
Average
80.6
77.0
53.8
75,3
12.9
52.6
81.8
76.2
66.2
79.1
82.4
89.4
91.2
Adroit Tasks
CQL.
IQL
BCQ
BEAR
O-RL
BC
DD
D-QL
DT
StAR
GDVT
QT
RVDT
pen human
37.5
71.5
66.9
-1.0
90.7
63.9
66.7
72.8
79.5
77.9
92.5
111.9
127.2
5.5
hammer-human
4.4
1.4
0.9
0.3
0.2
1.2
1.9
0.2
3.7
3.7
5.5
10.4
24.0
+ 1.5
pea-eloned
39.2
37.3
50.9
26.5
37.0
42.8
57.3
75.8
33.1
86.2
85.8
117.8
8.6
hammer--loned
2.1
2.1
0.4
0.3
2.0
0.6
1.7
3.1
3.0
0.3
8.9
11.8
21.3 ::: 2.7
Average
20.8
28.1
29.8
6.5
38.2
25.7
28.3
33.4
40.5
28.8
48.3
55.0
72.6
Kitchen Tasks
IQL
BCQ
BEAR
BC
D-QL
QT
RVDT
kitehen-Comp.
43.8
62.5
8.1
0.0
2.0
65.0
65.0
84.0
50.8
40.8
43.8
81.7
84.5 + 2.3
kitchen-parial
49.8
46.3
18.9
13.1
35.5
33.8
57.0
60.5
57.9
12.3
73.3
72.5
75.0 + 2.5
Average
46.8
54.4
13.5
6.6
18.8
49.4
61.0
72.2
54.4
26.6
58.6
77.1
79.8
Maze2D Tasks
CQL
IQL
BCQ
BEAR
TD3+BC
BC
Dittuser
DD
GDT
QDT
QT
RVDT
mazs2d-u
94.7
42.1
49.1
65.7
14,8
88.9
113.9
116.2
31.0
50.4
57.3
99.2
145.1 + 3.8
maz:2d-m
34.9
17.1
25.0
38.3
121.5
122.3
168.8
maze2d-1
49.6
61.7
30.8
81.0
88.6
1.5
123.0
125.9
2.3
0.7
31.0
242.7
254.3
1: 4.6
Average
62.0
46.2
32.3
57.2
55.2
42.9
119.5
121.5
13.8
19.6
33.9
170.2
194.3
AntMaze Tasks
CQL
IQL
BCQ
BEAR
TD3+BC
BC
D-QL
QT
RYDT
antmaze-u
74.0
87.5
78.9
73.0
78.6
54.6
73.1
93.4
59.2
51.3
76.0
96.0
98.0 + 4.0
antmaze-u-d
84.0
62.2
55.0
61.0
71.4
45.6
49.2
66.2
53.0
45.6
69.0
92.0
98.0 :: 4.0
antmare m-d
53.7
70.0
0.0
8.0
3.0
0.0
24.6
78.6
0.0
0.0
6.0
24.0
30.0
1. 6.3
antmaze-1-d
14.9
47.5
2.2
0.0
0.0
56.6
10.0
10.0 + 0.0
Average
34.0
35.5
73.7
57
Gym MuJoCo tasks are widely employed to evaluate continuous control capabilities under dense reward conditions. From Table MI we observe that RVDT achieves new SOTA results across all datasets. Notably, scores exceeding 100 on D4RL MuloCo tasks already near the theoretical performance ceiling, where boundary effects make even small improvements significant. Moreover, algorithms employing return-coverage rebalancing mechanisms, such RVDT, QT, CGDT, and TD3+BC, consistently demonstrate superior performance over traditional methods (e.g., BC and DT) across the majority of tasks, providing empirical support for our theoretical analysis.
Adroit tasks present substantial challenges for offline RL due to pronounced extrapolation errors stemming from the limited coverage of human demonstrations [[ Consequently, these tasks crit- ically test robustness against distributional shift and the ability to imitate from sparse expert data. Several observations emerge from our results: (i) except for RVDT and QT, other baselines fail to achieve satisfactory performance; (ii) RVDT and QT significantly outperform DT due to their utilization of rebalancing mechanisms; (iii) RVDT notably surpasses QT with 25% and 105% im- provements on pen and hammer, respectively. These findings underscore RVDT's effectiveness in learning from limited data and highlight the advantage gained from integrating explicit and implicit rebalancing mechanisms, effectively mitigating extrapolation errors and distributional shift problem.
Kitchen tasks primarily assess algorithms' capabilities in multi-task manipulation and generaliza- tion to unseen states. Our results indicate that RVDT slightly outperforms current SOTA methods, demonstrating its competence in managing multi-task dependencies and generalizing effectively. Furthermore, RVDT substantially surpasses non-CSM approaches, reinforcing the superiority of CSM methods over DP-based methods, aligning with findings from prior works [[]] [
Maze tasks evaluate the trajectory stitching capacity of CSM algorithms [[G. 225]. By integrating ex- pert KL regularization and Q-value guidance into Transformer- based policy, RVDT substantially
vii

---

Table 2: Performance comparison under varying proportions of expert trajectory augmentation. Each setting (e.g., 15% RVDT) corresponds to training on a mixed dataset composed of the base dataset (merged -m and -m-r) and top-k% expert trajectories sampled from the -m-e dataset.
0% Expert
15% Expert
30% Expert
50% Expert
Task
BC
QT
RVDT
BC
QT
RVDT
BC
QT
RVDT
BC
QT
RVDT
halfcheetah
42.4
46.9
49.3
46.0
79.5
92.7
73.3
91.3
93.3
71.1
92.8
94,4
hopper
65.4
95.8
100.2
99.0
108.0
113.1
112.3
110.6
113.8
112.8
108.3
113.9
walker2d
81.0
92.3
92.7
107.8
106.3
110.3
108.2
108.0
112.4
108.1
113.6
113.8
Average
62.9
78.3
80.3
84.3
97.9
105.4
97.9
103.4
106.5
97.3
104.9
107.0
Table 3: Performance comparison in low-data regimes. The learning difficulty increases progres- sively from D to 44, with D2/D2 containing 280 trajectories each and D3/D, containing 140 each.
D1
D2
D3
D4
Task (sparse R)
QT
RVDT
QT
RVDT
QT
RVDT
QT
RVDT
maze2d-umaze
100.3
171.7
81.8
101.8
73.1
76.9
61.4
100.5
mazc2d-medium
137.1
187.4
175.2
190.0
163.2
182.0
98.3
175.3
maze2d-large
109.5
140.4
81.5
90.1
104.4
131.3
100.2
95.1
Average
115.6
166.5
112.8
127.3
113.6
130.1
86.6
123.6
enhances its ability assemble optimal paths from fragmented sub-trajectories, thereby signifi- cantly outperforming all baselines. The AntMaze tasks, characterized by abundant suboptimal and unsuccessful trajectories, pose even greater challenges due to multi-action per state scenarios. While RVDT maintains competitive performance on smaller mazes (antmaze-u) virtually all CSM and value-regularized SL methods fail in larger, more complex mazes (antmaze-n, antmaze-1). In con- trast, D-QL, leveraging expressive diffusion policies, achieves the best performance on these chal- lenging tasks. This suggests that highly expressive policies are particularly effective in multi-action per state environments, highlighting the limitations in expressiveness of current CSM methods.
5.2 Performance under Imbalanced Trajectory Distributions
Effectiveness of Expert Data Utilization. We evaluate whether RVDT, integrating both explicit and implicit return rebalancing mechanisms, can more effectively leverage expert-level data from suboptimal datasets compared to QT and BC. We construct a base dataset by merging the -m and -m-r datasets, incrementally augmenting it with top-performing trajectories from the -m-e dataset. Table Z shows the performance of BC, QT, and RVDT at varying expert ratios. The results indicate: (i) RVDT consistently outperforms QT and BC at all expert ratios, highlighting the advantage of combining explicit and implicit rebalancing mechanisms compared to using only one (QT) or none (BC) of them; (ii) as the proportion of expert trajectories increases, RVDT and QT show steady improvements, while BC plateaus beyond approximately 30% expert data. This underscores the effectiveness of rebalancing mechanisms in enabling CSM policies to effectively utilize expert data.
Performance in Low-Data Regimes. To assess RVDT's performance in low-data scenarios [166), we construct four datasets (D)) by subsampling the original maze2d datasets. The learning diffi- culty increases progressively from DD to D4, with D; and D2 containing 280 trajectories each, and D3 and D4 containing 140 trajectories. Within each pair, D2 includes fewer expert-level trajectories than D1. as does DA compared to D3. Compared to the original 2000-trajectory maze2d dataset, these subsets better reflect algorithmic performance in small-sampi conditions (details for D provided in Appendix E.I. Table E shows RVDT consistently outperforming QT, often by large margin, with average improvements of 45.4%, 14.5% 14.2%, and 45.6% across Di-Da. This substantial performance gain highlights RVDT's superior ability to utilize useful information from limited and low-quality datasets (detailed comparisons in Appendix E.I..
Stitching Ability. Maze2D is commonly used to evaluate the trajectory stitching capability of of- fline RI. algorithms [25]], as agents are required to compose successful trajectories by integrating fragmented experience. We test RVDT and several baselines across four Maze2D environments
viii

---

Table 4: Performance comparison on Maze2D environments with sparse and dense rewards.
Dataset
CQL
DT
QDT
QT
RVDT
Sparse R
maze2d-open-vO
216.7 + 80.7
196.4
39.6
190.1 + 37.8
497.9
12.3
634.6 +
maze2d-umaze-vl
94.7 + 23.1
31.0
21.3
57.3 8.2
105.4 + 4.8
145.1 + 3.8
mazz2d-medhum--v
41.8 + 13.6
8.2
13.3
5.6
172.0 6.2
183.5 + 4,5
mazeed-large--I
49.6 8.4
2.3 0.9
31.0 19.8
240.1 2.5
254.3 + 4.6
Dense R
maze2d-open-vo
307.6 43.5
346.2 + 14.3
325.7
608.4 + 1.9
663.9 = 15.9
mazc2d-umaze-vl
72.7 10.1
-6.8 + 10.9
5.3
103.1 7.8
99.5 4.3
mazz2d-medium-vl
70.9 70,9 9.2
31.5 3.7
42.3 = 7.1
111.9 + 1.9
126.9 8.7
maze2d-large-v1
90.9 1 19.4
45.3 + 11.2
62.2 1 9.9
177.2 7.8
197.9 1 2.0
Table 5: Component-level breakdown across ablation variants.
Component
DT
DT-Dup RDT
QT
VDT
RVDT-Dup RVDT-Determ RVDT
Explicit Rebal. None
Dup.
KL
None
None
Dup.
KL
KL
Implicit Rebal. None
None
None
Q-value
Q-value
Q-value
Q-value Determ.
Q-value
Policy Type Determ.
Stoch.
Stoch.
Determ.
Stoch.
Stoch.
Stoch.
Table 6: Performance comparison of ablation variants across MuJoCo tasks.
Task
DT
DT-Dup
RDT
QT
VDT
RVDT-Dup RVDT-Determ
RVDT
halfcheetah
84.2
90.3
90.5
91.2
89.5
93.4
91.5
94.9
hopper
109.5
112.1
111.9
112.3
112.6
112.1
113.6
113.8
walker2d
108.2
108.8
109.7
113.2
110.3
110.9
113.1
118.7
Average
100.6
103.7
104.0
105.6
104.1
105.5
106.1
109.1
with increasing complexity (open, umaze, medium, large) under both sparse and dense reward settings. Results in Table gl show RVDT significantly outperforming baselines across almost all configurations. Notably, RVDT achieves significant improvements over QT under the sparse reward setting, where effective trajectory stitching is crucial. Results demonstrate that integrating explicit and implicit return rebalancing mechanisms substantially improves CSM-based policies' stitching capabilities, particularly for tasks requiring long-horizon credit assignment ability.
5.3 Ablation Studies
To quantify the contribution of each individual component in RVDT, we perform ablation studies on merged datasets (-m-e, -m-r, and -m) for MuJoCo continuous control tasks. These merged datasets alleviate coverage limitations in state, action, and return spaces, providing clearer evalua- tion of component impacts without dataset-boundary biases [G. 68]. The ablation variants focus on three key aspects: explicit rebalancing mechanisms, implicit rebalancing mechanisms, and policy rypes. Explicit rebalancing mechanisms include expert-policy KL regularization and expert-level data duplication, while the implicit rebalancing mechanism corresponds to Q-value guidance. For policy types, we consider both Gaussian stochastic policies and deterministic policies. The ablation variants are summarized in Table & Specifically, the two major regularizers introduced in our frame- work, the policy regularization term and the value-based guidance term, correspond to the "Explicit Rebal." (Explicit Rebalancing Mechanism) and "Implicit Rebal." (Implicit Rebalancing Mechanism) abbreviated in Table G. respectively. The variant RDT denotes the ablated version of RVDT without the value-based guidance component, while VDT represents the version without the policy regu- larization term. For the ablated variant employing deterministic policies (RVDT-Determ), the KL regularization term is replaced with an MSE objective. In variants where data duplication is applied, the top 50% of trajectories are directly duplicated and appended to the original dataset. Details are provided in Appendix Q
Table summarizes the ablation results and reveals several key insights: (i) RVDT consistently achieves top performance across all tasks, demonstrating the synergistic effectiveness of integrating all proposed components. (ii) Both RDT (explicit rebalancing) and QT/VDT (implicit rebalanc- ing) individually outperform DT, confirming the effectiveness of applying return-overrge rebalanc-
ix

---

ing mechanisms on DT, and indicating that the combination of rebalancing mechanisms in RVDT provides additional performance gains. (iii) RVDT-Dup, which employs direct expert-level data duplication, achieves comparable or superior performance relative to VDT, supporting our theoreti- cal insights regarding the benefits of explicit rebalancing mechanisms. (iv) comparison between RVDT-Determ (deterministic policy) and RVDT (stochastic policy) shows that RVDT with stochas- tic policy outperforms its deterministic counterpart, highlighting the importance of policy expres- siveness, particularly in multi-action per state scenarios.
6
Conclusion
This work establishes that both expert-level and full-spectrum return-coverage are critical determi- nants of the performance and sample efficiency of CSM policies in offline RL. Theoretical bounds that quantify the impact of these two aspects of return-coverage on policy optimality and sample complexity derived, offering a principled understanding of the underlying mechanisms that drive the effectiveness of return rebalancing strategies in CSM-based methods. Moreover, we propose RVDT, the first CSM-based approach that elegantly integrates explicit and implicit return rebalanc- ing mechanisms to effectively alleviate challenges posed by suboptimal datasets. Empirical evalua- tions on D4RL benchmarks demonstrate that RVDT enhances state-of of-te-of-thearrr performance in offline RL, validating the benefits and efficacy of return-coverage rebalancing.
Acknowledgements
This work was supported in part by the National Natural Science Foundation of China: 62206248.
References
[1] Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International conference on machine learning, pages 22-31. PMLR, 2017.
[2] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 2022.
[3] Saurabh Arora and Prashant Doshi. A survey of inverse reinforcement learning: Challenges, methods and progress. Artificial Intelligence, 297:103500, 2021.
[4] Wensong Bai, Chao Zhang, Yichao Fu, Peilin Zhao, and Hui Qian. Pacer: A fully push- forward-based distributional reinforcement learning algorithm. Neurocomputing, page 131301, 2025.
[5] Mare G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforce- ment learning. In International conference on machine learning, pages 449 458. PMLR, 2017.
[6] David Brandfonbrener, Alberto Bietti, Jacoh Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 15:1542-1553, 2022.
[7] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl with- out off-policy evaluation. Advances in neural information processing systems, 0144933-4443
2021.
[8] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language mod- els are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
[9] Yuanying Cai, Chuheng Zhang, Zhao, Wei Shen, Xuyun Zhang, Lei Song, Jiang Bian, Tao Qin, and Tieyan Liu. Td3 with reverse kl regularizer for offline reinforcement learning from mixed datasets. In 2022 IEEE International Conference on Data Mining (ICDM), pages 21-30. IEEE, 2022.
X

---

[10] Jiayu Chen, Bhargav Ganguly, Yang Xu, Yongsheng Mei, Tian Lan, and Vaneet Aggarwal. Deep generative models for offline policy learning: Tutorial, survey, and perspectives on future directions. Transactions on Machine Learning Research, 2024.
[11] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 54:15584-15097, 2021.
[12] Coleman, C Yeh, S Mussmann, B Mirzasoleiman, Bailis, Liang, J Leskovee, and M Za- haria. Selection via proxy: Efficient data selection for deep learning. In International Confer- ence on Learning Representations (ICLR), 2020.
[13] Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International conference on machine learning, pages 1096-1105. PMLR, 2018.
[14] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is es- sential for offline rl via supervised learning? In 10th International Conference on Learning Representations, ICLR 2022, 2022.
[15] Rasool Fakoor, Jonas W Mueller, Kavosh Asadi, Pratik Chaudhari, and Alexander J Smola. Continuous doubly constrained batch reinforcement learning. Advances Neural Information Processing Systems, 34: 1e111126-11272 2021.
[16] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072199 2020.
[17] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learn- ing. Advances in neural information processing systems, 34:20132-200 2021.
[18] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages 1587-1596. PMLR, 2018
[19] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning with- out exploration. In International conference on machine learning, pages 2052-2062 PMLR, 2019.
[20] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
[21] Zhang-We Wei Hong, Pulkit Agrawal, Remi Tachet des Combes, and Romain Laroche. Hamess-- ing mixed offline reinforcement learning datasets via trajectory weighting. In The Eleventh International Conference on Learning Representations.
[22] Zhang- Wei Hong, Tao Chen, Yen-Chen Lin, Joni Pajarinen, and Pulkit Agrawal. Topological experience replay. In International Conference on Learning Representations. OpenReview. net, 2022.
[23] Zhang- Wei Hong, Aviral Kumar, Sathwik Karnik, Abhishek Bhandwaldar, Akash Srivastava, Joni Pajarinen, Romain Laroche, Abhishek Gupta, and Pulkit Agrawal. Beyond uniform sam- pling: Offline reinforcement learning with imbalanced datasets. Advances Neural Informa- lion Processing Systems, 36:4985-5009, 2023.
[24] Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Has- selt, and David Silver. Distributed prioritized experience replay. In International Conference on Learning Representations, 2018
[25] Shengchao Hu, Ziqing Fan, Chaoqin Huang, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Q-value regularized transformer for offline reinforcement learning. International Con- ference on Machine Learning, pages 19165-19181. PMLR, 2024.
xi

---

[26] Shengchao Hu, Ziqing Fan, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Harmodt: Harmony multi-task decision transformer for offline reinforcement learning. In International Conference Machine Learning, pages 19182-19197. PMLR, 2024.
[27] Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Learning multi-agent communication from graph modeling perspective. In The Tweifth International Conference on Learning Rep- resentations.
[28] Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Graph decision transformer. arXiv preprint arXiv:2303.03744 2023.
[29] Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-tuning decision transformer with preference ranking. arXiv preprint arXiv:2305.0966 2023.
[30] Shengchao Hu, Shen, Ya Zhang, and Dacheng Tao. Graph decision transformer for offline reinforcement learning. SCIENCE CHINA-INFORMATION SCIENCES, 68(6), 2025.
[31] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pages 9902- 9915. PMLR, 2022.
[32] Vishal Kaushal, Rishabh Iyer, Suraj Kothawade, Rohan Mahadev, Khoshrav Doctor, and Ganesh Ramakrishnan. Learning from less data: A unified data subset selection and active learning framework for computer vision. 2019 IEEE Winter Conference on Applications of Computer Vision (WACV), pages 1289-1299. 2019.
[33] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing sys- lems, 33:21 182111100 810-21823, 2020.
[34] Sungyoon Kim, Yunseon Choi, Daiki E Matsunaga, and Kee-Eung Kim. Stitching sub- trajectories with conditional diffusion model for goal-conditioned offline rl. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 13160-13167. 2024.
[35] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6988 2014.
[36] Yilun Kong, Guozheng Ma, Qi Zhao, Haoyu Wang, Li Shen, Xueqian Wang, and Dacheng Tao. Mastering massive multi-task reinforcement learning via mixture-of-expert decision trans- former. In Forty-second International Conference on Machine Learning.
[37] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learnin. arXiv preprint arXiv:2110.06169, 2021.
[38] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Process- ing Systems, 32, 2019.
[39] Aviral Kumar, Joey Hong, Anikait Singh, and Sergey Levine. When should we prefer offline reinforcement learning over behavioral cloning? arXiv preprint arXiv:2204.056.11 2022.
[40] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 179- 1191, 2020.
[41] Thibault Lahire, Matthieu Geist, and Emmanuel Rachelson. Large batch experience replay. In International Conference Machine Learning, pages 1 11813. PMLR, 2022.
[42] Guanghe Li, Yixiang Shan, Zhengbang Zhu, Ting Long, and Weinan Zhang. Diffstitch: Boost- ing offline reinforcement learning with diffusion- trajectory stitching. In International Conference on Machine Learning, pages 28597-28608 PMLR, 2024.
[43] TP Lillicrap. Continuous control with deep reinforcement learning. arXiv:1509.02971, 2015.
arXiv preprint
xii

---

[44] Cong Lu, Philip Ball, Yee Whye Teh, and Jack Harker-Holder Synthetic experience replay. Advances in Neural Information Processing Systems, 36:463233-63444 2023.
[45] Jiafei Lyu, Xiu Li, and Zongqing Lu. Double check your state before trusting it: Confidence- aware bidirectional offline model-basee imagination. Advances in Neural Information Process- ing Systems, 35:38218 1-1333,, 2022.
[46] Michael Moret, Lukas Friedrich, Francesca Grisoni, Daniel Merk, and Gisbert Schneider, Gen- erative molecular design in low data regimes. Nature Machine Intelligence, 2(3):171-180, 2020.
[47] Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. Metadif- fuser: Diffusion model as conditional planner for offline meta-rl. In International Conference on Machine Learning, pages 26087-26105. PMLR, 2023.
[48] Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning a data diet: Finding important examples early in training. Advances in neural information processing systems, 34:20596-20007, 2021.
[49] Doina Precup, Richard S Sutton, and Satinder Singh. Eligibility traces for off- off-policy policy evaluation. In JCML, volume 2000, pages 759-766. Citeseer, 2000.
[50] Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithmss Cambridge university press, 2014.
[51] Jinghuan Shang, Kumara Kahatapitiya, Xiang Li, and Michael S Ryoo. Starformer: Trans- former with state-action-reaate representations for visual reinforcement learning. In European conference computer vision, pages 462-479. Springer. 2022.
[52] Yusuke Shinohara. A submodular optimization approach to sentence set selection. In 2014 IEEE International conference on acoustics, speech and signal processing (ICASSP), pages 4112-4115. IEEE, 2014.
[53] Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
[54] Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Ben- gio, and Geoffrey J Gordon. An empirical study of example forgetting during deep neural network learning. International Conference on Learning Representations, 2019.
[55] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[56] Jianhao Wang, Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang Offline reinforcement learning with reverse model-based imagination. Advances in Neural Information Processing Systems, 34:29420-29432, 2021.
[57] Yuanfu Wang, Chao Yang, Ying Wen, Yu Liu, and Yu Qiao. Critic-guided decision trans- former for offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 15706-15714, 2024.
[58] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusson policies expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations.
[59] Kai Wei, Rishabh Iyer, and Jeff Bilmes. Submodularity in data subset selection and active learning. In International conference on machine learning, pages 1954-1963. PMLR, 2015.
[60] Kai Wei, Yuzong Liu, Katrin Kirchhoff, and Jeff Bilmes. Unsupervised submodular subset selection for speech data. In 2014 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 4107-4111 IEEE, 2014.
[61] Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Flastic decision transformer. Advances in Neural Information Processing Systems, 36, 2024.
xiii

---

[62] Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodrigun Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline In Interna- lional Conference on Machine Learning, pages 18989-39009 PMIR, 2023.
[63] Teng Yan, Zhendong Ruan, Yaobang Cai, Yu Han, Wenxian Li, and Yang Zhang. Q- value regularized decision convformer for offline reinforcement learning. arXiv preprint arXiv:2409.0800 2024.
[64] Junjie Zhang, Jiafei Lyu, Xiaoteng Ma, Jiangpeng Yan, Jun Yang, Le Wan, and Xiu Li. Uncertainty-driven trajectory truncation for data augmentation in offline reinforcement learn- ing. In ECA/ 2023, pages 3018-3025. IOS Press, 2023.
[65] Hongling Zheng, Li Shen, Yong Luo, Tongliang Liu, Jialie Shen, and Dacheng Tao. Decom- posed prompt decision transformer for efficient unseen task generalization. Advances in Neural Information Processing Systems, 37:122984 123006, 2024.
[66] Hongling Zheng, Li Shen, Yong Luo, Deheng Ye, Bo Du, Jialie Shen, and Dacheng Tao. De- cision mixer: Integrating long-term and local dependencies via dynamic token selection for decision-making. In Forty-second International Conference on Machine Learning.
[67] Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. international conference machine learning, pages 27042-27059 PMLR, 2022.
[68] Zhaoyi Zhou, Chuning Zhu, Runlong Zhou, Qiwen Cui, Abhishek Gupta, and Simon Shaolei Du. Free from bellman completeness: Trajectory stitching via model-based return-conditioned supervised learning. In NeurIPS 2023 Foundation Models for Decision Making Workshop.
[69] Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Haoquan Guo, Tingt- ing Chen, and Weinan Zhang. Diffusion models for reinforcement learning: A survey. arXiv preprint arXiv:2311.01223, 2023.
xiv

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of work are discussed in experiments in section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
XV

---

Justification:
The complete set of assumptions is presented alongside the theoretical analy- sis in the main manuscript, and the corresponding proofs are provided in Appendix (
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The algorithmic details are thoroughly presented in the main manuscript, while the training and inference procedures are described in Appendix B Comprehensive experimental settings and configurations are provided in Appendix [
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or the contribution it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation,
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
xvi

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The paper utilizes open-source datasets, and the complete code of the algo- rithm is provided in the supplementary material, along with detailed instructions to faith- fully reproduce the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nniss.ce)
for more details.
While we encourage the release code and data, we understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (ittps:)
CoodeSuum ssionpolicy)
more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Section EL along with Appendix B and Appendix C provides a comprehensive description of the experimental settings and details, including all training and test configu- rations, data splits, hyperparameters, and optimizer selections.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The experimental results are presented with error bars that represent the stan- dard error of the mean (SEM), as reported in the corresponding tables (c.g., Table II).
Guidelines:
The answer NA means that the paper does not include experiments.
xvii

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The computational resources used for our experiments, including hardware specifications, memory usage, and execution time, are detailed Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: The research conducted in this paper fully complies with the NeurIPS Code of Ethics. All datasets used are publicly available and legally obtained, with no privacy concerns. The experimental evaluations are designed to ensure fairness and unbiased com- parisons across different methods. Additionally, the computational resources are managed responsibly, and the study does not pose any negative environmental or societal impacts.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
xviii

---

Answer: [NA]
Justification: This work primarily focuses on foundational research in machine learning. specifically in the development and evaluation of offline reinforcement learning algorithms. It does not involve applications or deployment scenarios that would directly impact societal, ethical, privacy-related domains. Therefore, the study does not present any identifiable positive or negative societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considcrations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This work focuses offline reinforcement learning algorithms evaluated on publicly available datasets (e.g.. D4RL benchmarks) that do not contain sensitive or scraped information. The proposed models are purely for experimental evaluation and do have direct deployment scenarios that would pose risks for misuse or dual-use concerns.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
xix

---

Justification:
All datasets, codebases, and pre pre-trained models used in this paper are properly cited, and their licenses and terms use are explicitly followed specified by the original authors.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswithcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification: This work introduces implementation of the RVDT algorithm, which is open-sourced and well-documented. The documentation provides sufficient instructions for reproduction and experimental setups, ensuring clarity and ease of use.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
XX

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsoureing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva-
lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non nta-standard components.
Please refer to our LLM policy /nttps://neur/ httpppppc///ooooeeeee////221//11 for what should or should not be described.
xxi

---

A Proofs
A.1 Proof of Theorem
The original proof is provided in [B]. For completeness, we restate the proof within our context to better align with the assumptions and notations adopted in this paper.
Proof. Let g (#1, (1:H) be the value of the return by rolling out the open loop sequence of actions a138 under the deterministic dynamics induced by T and Then we can write
Es, [S (81)) - (T²) E,, [Bxyles [S (81) sil]
=Es, [an 18: [S (81) g (81,)111)) + Rs, [oan [g (81, (1.)) - 9i]
S E,, [Ban [f (s1) (51, [1=]]] + eHf
(8)
where the last step follows by bounding the magnitude of the difference between Φ1 and (81, (1:H) by H and applying union bound over the H steps in the trajectory (using the near determinism assumption), namely:
H supup Par (rr # r (st, az) or 91+1 + T (st, ac) S eH?
Now we consider the first term from (S). Again bounding the magnitude of the difference by we get that
[
Es Bannwwyle [f (s1) 9 (81, [1H)] < Esg Jay
Pay (a1.
[1) (s), (1.) + (s1)
H
(9)
To simplify notation, let s, = T' (s1, d1:-1) be the result of following the deterministic dynamics defined by T up until step Expanding out, applying the near determinism, the consistency of f. the coverage assumption, canceling some terms, and then inducting we see that:
Pay (anerl (81) = nz (a1l01) /
P ((82(s1, a1) Pay (a2.0181,82)
S <yy (01(s)) Pr, (a2:1 81, 82) te
Pas (g) f (81) (81,ar)
(a1lss))
Pxy (a2.H(s1.22) +e
Px. (0) (s1) (s1)
+ n2 (g) r (81, = (81)
(81, 82) Px, (a2.1 81, 82) c
S STg (a1181)
PTs (g1 = f (81) (1))
c+ PT (g2 = f (2)) (22)
= (1)ls))
Pa, (02:1 32) te
Pso (gi f (s1 (11)
Pas (g2 (32) (32)
1
S ny (ails1)
Pay (02.1s1,, 52) +e
-++1
+1
Pap (g) f (81) (81)
aff
Pao (92 (52) (2)
Pro (gs (8g)
13, az)
I no (an (11 ne (02)32)
Pa, (as. 83)
Pa, (g) f (s1) (1)
PI, (0z = f (52) (22)
1
+ 2:
+1
af
H
TB (3))
Pa. (9H = f (5H) 131
He
- +1
1
Pxp (g) f (81) 181)
af
t=1
H
[[g (81, (1.H) = [
1 =+1)
my (a. (3:)
+ He
Pa, (g) = f (81) (81)
af
1=1
xxii

---

where the last step follows from the determinism of the trajectory that determines &μ and the consis- tency of f. Plugging this back into (H) and noticing that the two indicator functions can never both be 1, we get that:
(1) as
Bi Palarayla (x1) (s1,a):))]] He
Plugging this back into ( yields the result.
A.2 Proof of Theorem (performance gap)
We first give the following Lemmas help the proof of Theorem 0
Lemma (Brandfonbrener al. [G]). Consider the assumptions of Theorem (1 There exists opti- mal conditioning function f", defined such that (*) corresponds to the return of the deterministic optimal policy πt for the undertying MDP. Assume optimal return-coverage Pnς. (g f (s1))ss) 2
o'j holds for all initial states 81. Then we have that
- - < ( Oxf.
+3 2²e.
Proof. Since f*1) corresponds the return of the deterministic optimal policy x² for the under- lying MDP, so that E If" (s1]] is approximately J (=*).
Let T ($1, t) represent the state reached by running nr from $1 for t steps under the deterministic dynamics defined by T. Then:
H
= Σ (T" (81, (81, (o1.
fr (81)
1=1
Now we have as in the proof of Theorem that the probability that * (s) is bounded by eH, so that
Es, [f" J JIx' E,, Eg lss [f" (s1) all s E,, [Px (g (81) 81) H S cH²
Combining this with Theorem M yields the result.
Lemma 2 (Achiam al. [M]). Let de refer the marginal distribution of Px over states only. For any two policies nt we have:
dd d 11 2H Eand, [TV (F(-1s)]|=(( [s)]]]
Proof. First we will define few useful objects. Let (s) = Pa (sh s). Let A. = (s) d (s) Let 5n = 2 [TV (a(- l/ s))]
Now we claim that A. < 8n-1 A for and A =.
To see this, consider some fixed h. Note that (s) = S&x dh (s') Sar x (a' |s') Ps') P(s 8'aa'). Then expanding the definitions and adding and subtracting we see that
/ (db(s) (s)
A
SII dt (5)) / (I (T 89 of <' 8() ((sss'a))
(dh (8)) - d (6') / <'(a' 80) P (s 1''
<2R -2 TV (T(( ()(( s)) ||d - - 11, 5n A
xxiii

---

Now applying the claim and the definition of d we get that
H
H
4-1
H
|d dx 12 1
Σ An
S
Σ
5n = 2H Esnd, [TV (x(- |s]|=(( |s)]]
H
H
H
Now we are ready to prove Theorem
M.
Proof. Let (s1) denote the return of the optimal deterministic policy n² under deterministic dy-
namics. According to Theorem I and Lemma MI we have:
1_ +3 H²a
J(x') J(T Se E
(10)
ar..
| J( S [J( J(TT)] + +J((TTTMM))) - [(T))]-
(11)
Denote of" with an arbitrary conditioning function f' by Ty. LetP's x, and Px denote the station- ary distribution over states, actions, and returns induced by policy Ty and n3 respectively.
[(xp) -(TT)) = H|| (EP, [r(s, -Ep,, [r(s,a)]] (
< H dap -dxa dxp 11
$ 2112 Bendrs
[TV (T (T
U/ [xy(abb) [
< 2 Egndag
(12)
J. TT(aas)
Pap (r)))()) - Pap (s(()))))))
H H
Px. (f'ss)ss
2
V. To(als))
Egnddy
where the third inequality above utilizes the result of Lemma Denote Wanda, U. [a(as) by ag. By equation (I2) and equation (I have
[a J( (SMM) s 1 + aj
1 Hzas
(13)
Combine (IU) and (13 we conclude
1 _ +3)H² +
1 +(- H²
J(x²) - J( CSM S c(
(14)
oj
aj aj
Where we have the condition
a + aa. S 1,, (0,1);
(15)
if f 1).
(16)
That is 3 should have both good coverage on sample f and the optimal f*, when expert data increase, f will increase but f may decrease.
Should balance expert data and full- spectrum data!
(17)
xxiv

---

A.3 Proof of Theorem 3 (sample complexity)
We first introduce the following useful lemma, whose proof synthesizes the techniques from Achiam et al. [I]] and Brandfonbrener et al.[6]]
Lemma 3 (Sample complexity of CSM), Consider any conditioning function f. + R such
that the following assumptions hold: (i) Bounded occupancy mismatch: Paçau (s) S Cy (s)
for all s. (ii) Return-coverage: PRa(0
f(s) s) 2 af for all s. (iii) The policy class 11 is finite. (iv) log T(a | s.g) - log т (a' s',g')l S for any (a, (.,,. 9. s''''' and all π E II (v) The
approximation error of II is bounded by Capprox L.e, mingen ((T:) S Capprox.
Define the expected loss for CSM algorithm as
L(f)= Eg-Pes EogPooo (16) [L (Paa (-,1s-119911 (, -,,1))((1,,)))))))))
Then with probability at least 1 8,
Ci H2 Ve as
log [1888) 1/4
J ( J (+t) 0
+ Cappprox
(18)
Proof. Applying the definition of J and Lemma D we get
J (TT J (ft) = H Ep (EP.quulii (r(s, a)] - -..,,,, [1p))-E,,,,,, ((,,,,((,00)))
< H daaaa ddaga-
$ 2 Rood [TV (MFS () s))] H²
Expanding definitions, using the multiply and divide trick, and applying the assumptions:
2 Badd gom
[V ( C.M ()))][(((( s))]
S. |a - #(a ((lls,(()))
Eynd, .po
(f 8)
Pa (f(()) 5) / Paz(al s, f(s)) - #(a s. f.)((s)]]
Cj Bynday
Pag (f() / ( s,,(s) #(a f()))
0f
Cf Bandang
/ Prola 8) / |Pa (.9) #(a #
0f
C1 Eggnd Pandd
TV (1 (1.9)) [s,))]]
CJ (2)
af
where the last step comes from Pinsker's 's inequality. Combining with the above bound on the difference in expected values, we can get:
C 22(().
J (TT CSM) ((t)) < af
(19)
= We now write L(T) = [(T)- - Hmg where Hrp- = - -(.....)~Pr [log Prp (a 8,9)] and [(x)
-B(s,agg>>P, [log (a 8,9)] is the cross-entropy loss.
Denoting o ain L(T)) we have
L(ft) = L(#) L (TT) L (TT) < <((() - L( + Eapprox
XXV

---

Denoting the empirical cross-entropy loss that is minimized by म, we may further decompose
I(#t) - L(R)) L(+) L(#) - (x1) + ( ((*)
- 2 sup -()) - L(T))
xell
Under the assumptions on bounded loss differences, we may bound this, e.g., using McDiarmid's inequality and union bound on II obtain the final result in [50 6].
Now we are ready to prove the theorem
Proof. By Lemma B have the following result that
Ci H2 (od log 111/6
1/4
+ V Capprox )
(Ff - J (ft) so
(20)
as
By Theorem 2 we have the following result that
1 +3) + ( +
1 + H²C.
- $ of
(21)
af aj
Thus, by combining the results from (20) and (2I),, conclude that
J( - (+tg) so (
log /
1/4
C
H² Ve
Capprox +
+
H²
af
N
af
A.4
Proof of Proposition 11
Proof. To prove Proposition II, we show that the KL regularization term in (e) proportional the negative log-likelihone (NLL) loss, under the assumption that the policy is factorized Gaussian distribution with fixed isotropic covariance. The equivalence then implies that KL regularization explicitly reweights the NLL loss, thereby implementing explicit resampling mechanism.
Recall the KL divergence between multivariate Gaussians M(m, Σ1) and Wu Σ2) is given by:
T2 DKL W(H1,E)) (I Σx)) = tr("z' E1) + (#2 - -11 Σσ (122 - 11) - d log
Σ2
Σ
222)
where d is the dimensionality of the action space.
In setting, the policy is parameterized (a(a))) (s)(()), and the target is Dirac approximation centered at the action i.e., A(a, Substituting into (22), we obtain:
DKL (190 (a ET) = [ir(T)+ + (mo(s) - ((T)
(mo(s) - a) + log 1]
where we use the fact that Σ₁ = Eg = ET. Simplifying yields:
1
DKL ())
2 (1)) all2.
Taking expectation over (8, a) D gives:
Esaay DKL
((. = E [l]ma(s) all3].
2
which shows that the KI. divergence proportional the mean squared error (MSE) loss between (s) and the ground-truth action a.
xxvi

---

The log- ligeliiel of the Gaussian policy me(als) = No (o() {T) is:
1
log ( = log(2
lla (o(s)]??
25
Taking the negative log-lik oog-likeihood gives:
log (o(())) 25 |l - aal2 log(2
The constant term 2 (2 is independent of the policy parameters and does not affect optimiza- tion. Therefore, minimizing the KL divergence is equivalent (up scaling and additive constants) to minimizing the NLL loss:
(s.a)D DRI. (N(ua(s), [a, {!)) = E |- log Ta (a [s]s)] + constant.
(23)
Recall that the RDT objective in augments the standard NLL loss with a KL regularization term over expert trajectories:
H1
1 Σ
CRDT( ETWD
Σ
log (()) + to E, E.
PKL (To(-ss, ((o())s.gTT)) =
(1s:))
Te is delta function centered at a, (i.e., using the empirical action from the expert trajectory), and both policies are Gaussians with identical covariance çT, the KL term simplifies to an MSE loss, and hence to an NLL term by (25)). Thus, the second expectation term becomes equivalent to an additional NLL loss over De with weighting factor a.
This yields the final objective:
(R) = EInD (L+a. E De))
H
Σ - log To (a: g(T.) + constant.
Since the constant term does not affect the optimization process, the proof is thus completed.
A.5 Formal Statement and Proof of Proposition
Proof. To formally state Proposition E we first define the loss functions of the respective algorithms.
We begin by explicitly stating our key assumptions:
Infinite dataset size, i.e., DD + xo
Both π and ng are gaussian policy
Under these conditions, has been shown that minimizing the negative log- ligeliiel is equivalent to minimizing the KL divergence between and T..
NLL Dog
(24)
Eg-da g (-) (T(
As shown in Proposition II. KL regularization corresponds to a reweighting of the data distribution. We define the reweighted dataset D and the corresponding mixture policy zğ as:
1
D
Dr
1
1
(25)
1
T3
1 nB
me
With these definitions, we rewrite the learning objectives for DT, QT, RDT and RYDT:
म : arg max (D)(T) arg max E Endd, (T))) n
(26)
xxvii

---

के = arg max Lor(T) = arg max E Band,
+ 1) (=, a)]
(27)
Th == arg max (RDT(TT = arg max E B (tr.e//ndd,nl---Do((((())lll1l1 = T (28)
3 arg max LR) arg max E Banday D ( + ((0 (1)) (5. a)]]
(29)
Based on these optimization objectives, we can formally restate Proposition
Proposition 3 (Formal statement of Proposition D). Let gh. The Tg. and F denote the optimal poli- cies learned from the objectives of RVDT (29), RDT (28), QT (21) and DT (2t) respectively. Given the behavior policy иs there exists an expert-level subset De C D and an appropriate coefficient 2 0 such that the following inequalities hold for all states ~ daz.
E. (s) E, V (s) Es and
(30)
Eg V (s) 2 E, 2 E.V*a(s) (s).
Since Th T3. aim to prove the following policy improvement results:
E Bxndin (1) (FD S Exnd, ()) V=
(31)
E Egndeg Vo)(s) Egnday (s)
(32)
E Esnd, (1) (s) Esndy a (TT
(33)
E Fadd () V*) Bande (*)
(34)
Inequalities (GII) and (62) follow directly since (II) and (8)) special cases of (g) and (57). respec- tively, when a = 0.
We now prove 33
1/ Bedd ( (s))
B/and gggl (11)) QTT (1,
2 BBand
Q (s, a)] - E Eand, a00 (Dz (Dxz(a((ss))
(35)
2 1/ B/ -gg
1 (s, a)] - Esnd g (Tx((ss()s)
nB (
QQ (s, a)] nB ( (s)
where the first inequality follows the non-negativity of KL divergence, the second inequality follows Tq is the optimal policy of 27
Similarly, for (64)):
n B (s)()))
BBadd t Q (s,a)
(36)
2 B (1n)
Q _ (s,a)] - E Exnd, (
n/ BBaadd a
Q " a)) ()Farda (1) () (s))
This completes the proof.
xxviii

---

B Training and Inference
We instantiate the previously described formulation using a transformer-based architecture. The RVDT architecture extends the DT framework by explicitly accounting for stochastic policies, pre- dicting the policy's 's mean and log-variance through two separate fully connected layers. The com- prehensive training procedure for RVDT summarized in Algorithm M. For notational simplicity, Algorithm M represents the policy input using the shorthand (- Kiú, 0i-Kxi-1, 9i- ()-K)), which explicitly treats the current state-return pair (8t; 9:) as part of the context window. When i- K the context window is truncated to the available prefix of the trajectory.
In scenarios without supplementary datasets direct access expert policies, we rely on cumulative returns to differentiate trajectory quality. Specifically, we select the top p% trajectories (p is hyperparameterr) ranked by cumulative return, form the expert dataset De. The expert policy T" is subsequently trained via behavior cloning on De. Note that alternative imitation learning approaches such as GAIL [20] or IRL [3] could similarly be utilized derive T".
RVDT employs a two-step sampling process to ensure uniform sampling of sub-trajectories of length K" from the replay buffer Treplay- Initially, an entire trajectory is sampled with probability propor- tional to its length, and subsequently, a sub-trajectory of length K is uniformly extracted. For environments characterized by non-negative dense rewards, this approach resembles importance sampling scheme [62]].
We leverage offline Dynamic Programming techniques to train the value network Qp. optimizing via Temporal Difference (TD) learning [[:3]]_ In practical implementation, standard stabilization techniques such double Q-learning [|о]] and the use of target network | adopted. It is noteworthy that in the offline RL setting, conservative Q-learning (CQL) [+0] is typically employed to mitigate the out-of-distribution issue by maintaining policy similarity to the behavior policy. How- ever, since the RVDT inherently incorporates behavior cloning, which naturally ensures similarity to the behavior policy, the additional adoption of CQL yields minimal performance gains.
During inference, return-conditioned supervised methods initially require specifying a target return as the first RTG token, subsequently subtracting sampled rewards to update the RTG token iteratively, According to Theorem achieving the target return critically depends on the coverage of the speci- fied return in the dataset. Hence, a practical choice is to utilize the maximum return observed in the training dataset as the initial target. Nonetheless, given RVDT's integration with dynamic program- ming principles and its strong trajectory stitching capability, the actual optimal target return may surpass the dataset's 's maximum observed return. Building upon the inference strategy of QT, we ran- domly sample candidate returns within positive 10% margin above the maximum return observed in the dataset, thereby obtaining multiple candidate return-to-go tokens 98- 06......99 -0000 Actions corresponding each candidate return-to-oo token simultaneously generated. Subsequently, we employ the learned Q-value function to preferentially select actions associated with higher returns. This selection mechanism is formally expressed as: af arg maxa; Qw where az ~ x-118t, gi, =K 1)
C Experimental Details
C.1 Benchmark Details
Our evaluation performed across five diverse domains in the D4RL benchmarks [[ Gym, Adroit, Kitchen, Maze2D, and AntMaze. These domains are designed to comprehensively evalu- ate the capability of offline RL. algorithms in handling different types of dynamics, reward structures, and state-action distributions.
Gym-MuJoCo locomotion tasks serve as standard benchmarks for continuous control. These tasks, such as HalfCheetah, Hopper, and Walker2d, are characterized by smooth reward func- tions and dense feedback signals. The datasets include trajectories generated by policies of varying quality: random, medium, expert, and medium-experes representing different levels of optimal- ity. While expert and medium-expert provide high-quality trajectories, medium, random, and medium-replay contain suboptimal and exploratory behavior, offering diverse state-action cov- erage for offline learning.
Xxix

---

Algorithm 1: RVDT: Return-rebalanced Value-regularized Decision Transformer
Input:
Sequence length K, dataset 'D hyperparameters a; 1), expert ratio p%
Initialize: Policy To. critic Q
Construct expert dataset Dr by selecting top p% trajectories from D; Train expert policy nh on D. via BC;
for t to do
Sample trajectory 7 D and length- subsequence from
Denote the sampled subsequence as {(s,, ai, ) by re-indexing;
// Critic update
Update critic Qs using offline TD learning on transitions sampled from D;
// Policy update
for i to K do
Sample action a; ~ TT(-|si- Kús di Ki-K,, Ji ku)); Compute
ā:):
Compute KT (TT[ Ot-Ka-1, 9:- 9i-ks:) 1 (-1s:))
Update policy vy by minimizing Eq. (Z);
return :e:
Adroit tasks involve high-dimensional dexterous manipulation using a 24-DoF robotic hand. The datasets are primarily collected from human demonstrations and behavior cloned policies, result- ing in narrow state-action coverage and strong multi-modal behavior. Due to the inherent complex- ity of the hand's control space, policies must be effectively regularized to avoid out-of-distribution (OOD) actions. The sparse nature of successful demonstrations and the high variance in human trajectories further challenge offline RL algorithms in terms of generalization and stability.
Kitchen environment introduces mng-horizoon multi-task manipulation scenarios, where agent is required to execute a fixed sequence of four subtasks to achieve goal configuration, such as turning the microwave, opening the cabinet, switching on the light. This setting emphasizes the importance of temporal abstraction and multi-task planning. The datasets are collected from demonstrator policies that exhibit varying degrees of task completion, leading to fragmented sub- trajectories that need to be stitched together effectively during offline learning.
Maze2D tasks are specifically designed to evaluate the ability of offline RL algorithms to stitch to- gether fragmented sub-trajectories and recover globally optimal paths in continuous navigation setting. Unlike Gym-MuJoCo, Maze2D environments present sparse rewards, requiring the agent to combine distant state transitions to solve navigation tasks effectively. Although Maze2D also provides a dense reward setting, we primarily focus on its sparse reward version throughout this work unless otherwise specified. The agent controls 2D ball in complex maze-like structures, where the challenge lies in leveraging suboptimal paths to construct optimal global trajectories.
AntMaze represents a more challenging extension of Maze2D by introducing an 8-DoF And robot instead of simple 2D ball. The sparse reward structure remains, but the high-dimensional state space and unstable dynamics the Ant make trajectory stitching significantly more difficult. Suc- cessful navigation requires effective trajectory stitching and long-horizon credit assignment to traverse disconnected regions of the maze. Datasets are collected using goal-conditioned poli- cies with diverse start-goal configurations, introducing variability in trajectory quality and return distributions.
Dataset descriptions and policy types for representative tasks in the D4RL benchmarks is given in Table Ø
XXX

---

Table 7: Dataset descriptions and policy types for representative tasks in the D4RL benchmarks.
Task
Policy Type
Description
hupper-medium-cxpert
Expert + Medium policy
50-50 mixture trajectories from a near-optimal ex- pert policy and medium performance policy.
hupper-medium--2
Medium behavior policy
Collected by policy trained to partial convergence, rep-
resenting moderately optimal behavior.
"upper-nedduu--repaay-22
Medium Experience replay
A replay buffer from the medium policy, including
wide range of suboptimal and early-stage transitions.
pen-human-vl
Human demonstrations
Trajectories collected via human telcoperation using a virtual reality interface.
pen-cloned-vl
Behavior cloning policy
Trajectories generated by policy trained via behavioral cloning on human demonstration data.
kitchen-complate-v0
Export policy
Expert demonstrations accomplishing fixed sequence of four goal-directed tasks.
kitchen-partia-vv0
Expert policy
Expert demonstrations completing variable subset of kitchen tasks different orders.
maze2d-umaze-v
Trajectory tracking (suboptimal)
Random-wand trajectories collected small U-shaped maze.
maze2d-medium-v1
Trajectory tracking (suboptimal)
Random- walk trajectories a medium- sized maze increased spatial complexity.
maze2d-large-vl
Trajectory tracking (suhoptimal)
Random- walk trajectories the largest maze with long- horizon navigation.
antmaze-umaze-vO
Expert policy
Goal-conditioed expert demonstrations small shaped maze.
untmaze-umaze-divree---0
Expert policy (diverse goals)
Goal-rcaching demonstrations in the small maze with diverse start-goul pairs.
uttmaze-mediun--diccee--d
Expert policy (diverse goals)
Goal-directer expert trajectories in medium maze with varying start-goal configurations.
utmaze-large-diieree-xU
Expert policy (sparse diverse)
Sparse, long-horizon goal-reaching demonstrations in large maze using 8-DoF quadruped agent.
D Implementation Details
The policy network of RVDT is implemented as Decision Transformer, built upon the open-source minGPT codebase and the official QT implementation Detailed model parameters are provided in Table &
The expert policy x" is modeled using fully connected neural network and trained via behavior cloning. The corresponding model parameters are listed in Table g
The Q-networks represented by three-layer MLPs with Mish activations and 256 hidden units for each layer.
All networks are trained using the Adam optimizer
The code of RVDT, including all perparaameeerr configurations and experimental setups, is provided in the supplementary material for reproducibility.
The experiments were conducted sors and 8 NVIDIA GeForce RTX 4090 GPUs with 24 GB of memory. For each task, the average training time single GPU was approximately 20,000 seconds, with memory consumption of around 3000MB. These computational settings ensure reproducibility and align with the reported performance metrics.
server equipped with two AMD EPYC 7542 32-Core Proces-
htpp:://tthab.ccm/aarrahhhaaaaap https://gihuub.com/taaaaeeeeeee
XXxi

---

Table 8: Hyperparameters of RVDT in our experiment.
Parameter
Value
Number of layers
Number of attention heads
Embedding dimension
256
Nonlinearity function
ReLU
Batch size
256
Context length K
20
Dropout
0.1
Learning rate
3.0e
Weight decay
1.0e
Table 9: Hyperparameters of de in our experiment.
Parameter
Value
Number of layers
4
Embedding dimension
256
Nonlinearity function
ReLU
Batch size
256
Dropout
0.1
Learning rate
1.0e 4
Weight decay
1.0e 4
E Hyperparameters for RVDT
We recall the loss function for RVDT follows:
(RVDT(() =BB,~D E log ((:)),
+ A.END. Σnn KL no ((. ,/t 1| (-1s:)
(37)
E Bra
E R. (s) Qp (8,, [;)
(77)
involves four key hyperparameters that may vary across different tasks:
The weight of the NLL, A;;
The weight of the expert policy KL regularizer, a;
The weight of the Q-value guidance term, f);
The proportion of expert-level trajectories selected from the dataset D to form the expert dataset D., denoted p%.
We employ the Tree-structured Parzen Estimator (TPE) for hyperparameter optimization in each environment. The optimal values are determined based on performance on a held-out validation set, aiming achieve the best trade-off between policy regularization and value maximization. The key hyperparameters of RVDT given Table II Other hyperparameter configurations used for reproducing the experiments are provided in the code available in the supplementary material.
xxxii

---

Table 10: Key hyperparameter settings for all selected tasks to reproduce the main results reported Table
Tasks
a
6
Expert dataset ratio p% Grad Norm
halfcheetah- www.-mmeedee--e-------e
0.1
5.0
0.02
50%
15.0
hopper-medium- experttv22
0.01
1.0
0.02
50%
9.0
walker2d-medinee mediuu--pp-----
0.01
2.0
5c-6
10%
5.0
halfchectah-meding
0.1
5.0
0.02
10%
15.0
hopper-medium--2
1.0
1.0
0.02
50%
9.0
walker2d-mediuum-22
0.1
2.0
0.02
5.0
hallchhethh-mdddiu-eelaaayy2
1.0
5.0
2e-4
50%
15.0
hopper-meddium--eppayy222
1.0
3.0
0.02
50%
9.0
walker2d walker2d-medium--ree medium-replay--2
0.1
2.0
0.02
10%
5.0
maze2d-open-Vo
0.01
0.02
10%
maze2d-umaze-vl
1.0
5.0
1.0
30%
20.0
mazc2d-medium--rl
0.1
5.0
5e-6
30%
9.0
maze2d-large-vl
1.0
4.0
2c-6
30%
9.0
antmaze-umaze-vO
0.1
0.05
0.02
30%
9.0
autmazz-umazze-divere-----
0.1
0.02
0.01
70%
pen-human-vl
10.0
0.1
2e-3
50%
9.0
hammer-human-v/
1.0
1.0
0.02
50%
5.0
pen-eloned-vl
10.0
0.1
0.02
30%
9.0
hammer-cloned-vl
0.5
0.01
2c-4
50%
9.0
kitchen-come
circhen-complee-v00
0.1
0.001
1c-4
50%
9.0
kitchen-partial-v00
10.0
0.01
0.02
50%
9.0
maze2d-open-denne--00
1.0
0.01
0.02
10%
9.0
maze2d-umazz--ennsee.
0.5
3.0
0.02
10%
5.0
maze2d-mede mene2dd-mdiuumddesee
100.0
5.0
2e-4
maze2d-large-den
1.0
4.0
2e-3
10%
9.0
gutmaaz-mediimm-diveess----
0.1
0.01
0.02
30%
9.0
antmare-largeedivvree-vv
0.1
0.005
0.02
30%
9.0
Additional Experimental Results
F.1 Performance in Low-data Regimes
We provide a detailed description of the four sub-datasets D - D4 constructed for the low-data regime experiments. First, note that each of the original maze2d-umaze maze2d-large datasets consists of 2000 trajectories. For each task, we filter the original dataset to distinguish trajectories of different quality levels. We define three quality- based sub-datasets: expert, medium, and noise, denoted as DD., Dm. and D.. respectively. These subsets are obtained by randomly sampling 200 trajectories from the top 30%, middle 40%, and bottom 30% of returns in the original dataset.
maze2d-medine and
The sub-datasets D1 D, then constructed by randomly sampling from De. Dm, and D, ac- cording to the following proportions:
D1: 50% D. + 50% Dm + 40% D.,, comprising 280 trajectories, designed to achieve full return-coverage.
2. D2: 20% D. 20% Dm + 100% D., comprising 280 trajectories, dominated by random- policy data with a smaller fraction of medium and expert trajectories.
3. D3: 25% D. + 25% D,, + 20% D,,, comprising 140 trajectories, mirroring the empirical return distribution of D, but with fewer samples.
4. D4: 10% D_ + 10% D,, + 50% D,,, comprising 140 trajectories, mirroring the empirical return distribution DD with reduced sample size.
Table demonstrates the significant performance improvements RVDT over QT low-data regimes. The percentage improvements of RVDT over QT, derived from Table B are summarized in Table I. These results confirm that RVDT, by incorporating both explicit and implicit rebalancing mechanisms, substantially enhances its ability to recover optimal policies from limited data sam-
xxxiii

---

ples. This further highlights the non-trivial nature of RVDT's pioneering integration of explicit and implicit rebalancing strategies.
Table 11: Percentage improvement of RVDT over QT across different datasets. For example, an 71.2% indicates that RVDT outperforms QT by 71.2% on the maze2d-uma environment when trained with the D dataset.
Tasks (sparse R)
D1
D2
D3
D4
maze2d-umaze
71.2%
24.4%
5.2%
63.7% 1
maze2d-medium
36.7%
8.4% t
11.5% 1
78.3% 1
maze2d-large
28.2%
10.6%
25.8%
-5.1% 1
Average 45.4% 1
14.5%
14.2% 1
45.6% 1
F.2 Impact of Expert dataset ratio
We evaluate the performance of RVDT when trained with varying proportions of subsampled expert trajectories, i.e., by adjusting the expert dataset ratio p%. Adjusting the proportion p% directly affects the trade-off between az and a ; as the proportion increases from 10% to 100%, az first rises and then declines, while o exhibits the opposite trend. This experiment aims to (1) examine how the interplay between o and of influences algorithmic performance, and (2) identify the optimal ratio of expert-level trajectories required for effective learning.
The results, summarized in Table M.. show that using approximately 30% of the top- cop-perforning trajectories consistently yields the best results across Maze2D tasks. Specifically, the 30% expert subset achieves higher scores than both smaller (10%) and larger (50%, 70%, and 100%) proportions. suggesting that oj and o reach near-optimal balance under this configuration.
This observation aligns closely with our theoretical analysis, which highlights the importance of maintaining balance between expert-level and full-spectrum return-coverage It reveals that effec- tive rebalancing does not simply entail maximizing the proportion of high-quality expert trajectories. Intuitively, relying exclusively expert data (i.e., maximizing aj while neglecting ay) can cause significant performance degradation due to distributional shift, particularly when training data are limited. Therefore, careful calibration between expert-level and full-spectrum return-coverage crucial for achieving optimal performance in offline RL.
Table 12: Performance comparison with different proportions of subsampled expert datasets. For example, 10% RVDT corresponds to the performance RVDT when the expert dataset D. consists of the top 10% of trajectories ranked by return from the full offline dataset D.
Task (sparse R)
10% RVDT 30% RVDT
50%
RVDT
70% RVDT 100% RVDT
maze2d-umaze-vl
93.1
143.0
100.4
102.4
101.6
maze24-medium-vl
189.9
195.0
192.5
191.6
188.6
maze2d-large-vl
246.5
250.0
248.5
248.1
243.6
Average
176.50
196.0
180.47
180.70
177.93
F.3 Influence of Expert Data Quality
To empirically investigate how the quality of the expert policy affects RVDT performance, evalu- ate RVDT under varying proportions of expert trajectory augmentation. As summarized in Table IV each configuration (e.g., RVDT (15%)) is trained a mixed dataset composed of the base dataset (merged -m and -m-r) and the top-&% expert trajectories sampled from the -m-e dataset, which can be regarded expert-level trajectories.
The results show that RVDT performance consistently improves as the proportion expert trajecto- ries increases, with the most notable gain observed between RVDT (0%) and RVDT (15%). These findings indicate that higher-quality expert policies contribute to better performance. Combined with the results in Table [2 can clearly observe the influence of adjusting and rebalancing return
XXXiv

---

coverage on algorithm performance, thereby further confirming the effectiveness of the proposed cevern-cveerage rebalancing mechanism.
is important to note that the expert policy used in RVDT for the main experiments (see Table I4] is not trained from any externally provided demonstrations. Instead, it is derived from subset of the existing offline dataset, denoted as De C D, by selecting relatively high-return trajectories. Con- sequently, our method is fully applicable to standard offline RL settings without requiring any addi- tional expert data. Moreover, the return coverage assumed in our theoretical analysis is not restricted to strictly optimal trajectories. In practice, the presence of near-expert or moderately high-return tra- jectories is often sufficient to ensure effective learning, This property underscores that RVDT does rely on idealized expert-level datasets and remains robust as long as the dataset exhibits suffi- cient diversity to support relative return ranking. As demonstrated in Table [4. RVDT continues to achieve notable improvements even on the -medium and -medium-replay datasets, which contain no expert demonstrations, validating its practicality under realistic offline RL conditions.
Table 13: Performance comparison under varying proportions expert trajectory augmentation. Each configuration (e.g., RVDT (15%)) is trained on mixed dataset composed of the base dataset (merged -m and -m-r) and top-k% expert trajectories sampled from the -m-e dataset.
Task
RVDT (0%)
RVDT (15%)
RVDT (30%)
RVDT (50%)
halfcheetah
49.3
92.7
93.3
94.4
hopper
100.2
113.1
113.8
113.9
walker2d
92.7
110.3
112.4
113.8
Average
80.3
105.4
106.5
107.0
Table 14: Performance comparison across D4RL tasks. For each task, RVDT reports the mean and standard error of normalized scores [T66] over 30 random rollouts and averaged over 5 random seeds.
Gym Tasks
CQL
IQL
BCQ
TD3-BC
McRel
BC
DD
DT
SLAR
GDT
CGDT
QT
RVDT
halfcheetab mee
91.6
86.7
69.6
90.7
53.3
55.2
90.6
86.8
93.7
93.2
93.6
93.2 :: 0.8
94.4 1: 0.1
hopper-m-e
105
91.5
109.1
98.0
108.7
52.5
111.8
107.6
111.1
111.1
107.6
113.0 = 0.2
113.1
+ 0.5
wallker2d-m-e
67.3
107.5
108.1
109.3
112.0
112.7
hallcheetab-mn
49.2
47.4
41.5
48.4
42.1
42.6
49.1
42.6
42.9
42.9
43.0
51.0
0.2 0.2
$1.9 - 0.3
hopper-
65.1
67.6
96.9
100.2
walker2d-m
83.0
78.3
52.0
83.7
71.8
75.3
82.5
74.0
73.8
76.5
79.1
81.2 = 1.0
90.2 +
halfcheetab-ma.
45.5
44.2
34,8
44.6
36.6
39.3
36.6
36.8
40.5
40.4
48.8
1.5
$3.8 1: 2.0
hopper-m-r
95.0
94.7
31.1
60.9
93.6
18.1
100.0
82.7
29.2
85.3
93.4
102.1
103.2
walker2d-m-r
77.2
73.9
13.7
81.8
49.8
32.3
75.0
79.4
39.8
77.5
78.1
97.8 + 0.9
99.3 +: 0.8
Average
80.6
77.0
53.8
75.3
72.9
52.6
81.8
76.2
66.2
79.1
82/4
89.4
91.2
Adroit Tasks
CQI.
IQL.
BCQ
BEAR
O-RI.
BC
DD
D-QI.
TYT
S1AR
GDT
QT
RYDT
pea-hunan
37.5
71.5
66.9
-1.0
90.7
63.9
66.7
72.8
79.5
77.9
92.5
111.9
11.8
127.2 + 5.5
hammer-
0.9
1.2
3.7
5.5
10.4
5.3
pen-cloncd
39.2
37.5
50,9
26.5
60.0
37.0
42.8
57.3
75.8
33.1
86.2
85.8 + 12.6
117.8
8.6
hammer-eloned
2.1
0.4
0.6
1.7
3.0
8.9
11.8
5.8
21.3
20.8
28.1
29.8
6.5
38.2
25.7
28.3
33.4
40.5
28.8
48.3
58.0
72.6
Kiteben Tasks
IQL
BCQ
BC
D-QU
1YT
GDT
RVDT
kitchen-Come Combe
43.8
62.5
8.1
65.0
65.0
84.0
50.8
40.8
43.8
81.7
:: 1.7
84.5 4: 2.3 75.0 J: 2.5
kitchen-
kitchen-purtial partial
49.8
46.3
18.9
13.1
35.5
33.8
57.0
60.5
57.9
12.3
73.3
72.5 d: 2.5
Average
46.8
54.4
13.5
6.6
18.8
49.4
61.0
72.2
54.4
26.6
58.6
77.1
79.8
Maze2D Tasks
CQI.
IQ1.
BCQ
BEAR
TD3+BC
BC
Diffuser
DD
DYT
GDT
QDT
QT
RYDT
muze2d-u
94.7
42.1
49.1
65.7
14.8
88.9
113.9
116.2
31.0
50.4
$7.3
99.2 4: 1.0
145.1 3.8
manc2d-m
41.8
34,9
17.1
25.0
62.1
38.3
121.5
122.3
8.2
7.8
13.3
168.8
= 8.5
183.5
man:2d-1
49.6
61.7
30.8
81.0
88.6
1.5
123.0
125.9
2.3
31.0
242.7 = 7.3
254.3
Average
46.2
32.3
57.2
55.2
42.9
119.5
121.5
13.8
19.6
33.9
170.2
194.3
AntMaae Tasks
CQL
JQL
BCQ
BEAR
TD3-BC
BC
DD
D-QL
DT
SIAR
GDT
QT
RYDT
antmaze-4i
87.5
78.9
73.0
78.6
54.6
73.1
93.4
59.2
51.3
76.0
96.0
3.0
98.0 4.0
antmaze-ndd
84.0
62.2
55.0
61.0
71.4
45.6
49.2
66.2
53.0
45.6
69.0
92.0
1 6.
98.0 +: 4.0
antmaze
0.0
78.6
6.0
24.0
5.6
30.0
animaze-1-d
14.9
47.5
2.2
0.0
0.0
0.0
7.5
56.6
0.0
0.0
0.0
10.0 = 0.0
10.0 +
Average
56.6
66.8
34.0
35.5
38.2
25.0
38.6
73.7
28.0
24.2
37.8
57
59.0
XXXV

---

F.4 Details for Ablation Studies
The ablation variants focus on three key aspects: explicit rebalancing mechanisms, implicit rebal- ancing mechanisms, and policy types. Explicit rebalancing mechanisms include expert-policy KL regularization and expert-level data duplication. For the latter, the duplicated expert data correspond to the top 50% of trajectories ranked by return, with single additional copy included. The implicit rebalancing mechanism refers specifically to Q-value guidance. Regarding policy types, we consider two policy types, Gaussian stochastic policies and deterministic policies.
Table 15: Component-level breakdown across ablation variants.
Component
DT
DT-Dup RDT
QT
VDT
RVDT-Dup RVDT-Determ RVDT
Explicit Rebal.
None
Dup.
KL
None
None
Dup.
KL
KL
Implicit Rebal. None
None
None
Q-value
Q-value
Q-value
Q-value
Q-value Stoch.
Policy Type Determ.
Stoch.
Stoch.
Determ.
Stoch.
Stoch.
Determ.
Detailed descriptions of the ablation variants in Table
B are as follows:
DT: The original Decision Transformer method utilizing deterministic policies, implemented by disabling the Q-value guidance term of QT.
DT-Dup: stochastic variant of DT enhanced with expert data duplication, modeling the policy parameterized Gaussian distribution.
RDT: RVDT variant without the implicit rebalancing mechanism (Q-value guidance). isolating the effect of the expert-policy KI. regularizer.
QT: The original QT model implemented using the official codebase, serving the deterministic- policy counterpart of RVDT without explicit rebalancing mechanisms.
VDT: An RVDT variant excluding expert-polie KL regularization to specifically isolate the con- tribution of Q-value guidance.
RVDT-Dup: An RVDT variant employing expert-level data duplication as the explicit rebalanc- ing mechanism, enabling comparison between KL-based rebalancing and data duplication ap- proaches.
RVDT-Determ A deterministic policy variant of RVDT, replacing the stochastic policy with deterministic outputs and substituting KI. regularization with an MSE loss, assessing the role of policy expressiveness.
The complete RVDT method integrating expert-policy KL regularization and Q-value guidance with stochastic policy parameterized by Gaussian distribution.
RVDT:
Influence Policy Expressiveness on Performance
To further investigate the influence of policy expressiveness [ algorithmic performance, compare RYDT with several representative Q-learning-base based baselines (CQL, IQL, and D-QL) on the D4RL benchmark. As shown in Table [[ CQL and IQL employ Gaussian stochastic policies, while D-QL adopts a more expressive diffusion policy. In contrast, RVDT is built upon conditional sequence modeling (CSM) techniques and incorporates simple Gaussian stochastic policy.
From the results, we observe that D-QL, which leverages the diffusion policy, significantly outper- forms the simpler Gaussian-policy baselines (CQL and IQL), both across most tasks and in terms of overall average performance, demonstrating the advantage of higher policy expressiveness. Never- theless, RVDT achieves superior performance most tasks, which highlights the effectiveness of leveraging CSM together with the return- returncoverrge rebalancing mechanism.
However, in the AntMaze tasks with larger mazes (-m, -1), D-QL performs better than RVDT. This may attributed to the multi-legged morphology of the agent, which entails a complex and highly multi-modal action distribution. In these scenarios, the expressiveness of the diffusion policy becomes particularly beneficial, whereas RVDT's Gaussian policy may be insufficient capture such rich multimodal behaviors.
XXXVi

---

Table 16: Performance of Q-learning- based baselines (CQL, IQL, D-QL) and RVDT on D4RL tasks, where CQL and IQL adopt Gaussian policies, and D-QL employs a diffusion policy. Results are averaged over 5 random seeds.
Tasks
CQL
IQL
D-QL
RVDT
pen-human
37.5
71.5
72.8
127.2
hammer-human
4.4
1.4
0.2
24.0
pen-cloned
39.2
37.3
57.3
117.8
hammer-cloned
2.1
2.1
3.1
21.3
kitchen-Comp
43.8
62.5
84.0
84.5
kitchen-partial
49.8
46.3
60.5
75.0
maze2d-m
41.8
34.9
91.0
183.5
maze2d-1
49.6
61.7
200.7
254.3
antmaze-u
74.0
87.5
93.4
98.0
antmaze-u-d
84.0
62.2
66.2
98.0
antmaze-m-d
53.7
70.0
78.6
30.0
antmaze-l-d
14.9
47.5
56.6
10.0
Average
41.2
48.7
72.0
93.5
Overall, these findings suggest that both policy expressiveness and CSM with return-coverage rebal- ancing effective factors in improving offline RI. performance. The mechanisms comple- mentary: diffusion policies enhance local action diversity, while CSM with rebalancing strengthens long-horizon modeling capability. Exploring their integration represents promising direction for future research.
F.6
Training Time
To evaluate the computational efficiency of RVDT, we provide detailed comparison of training times among DT, QT, and RYDT on D4RL, benchmark. The results, summarized Table II report the total training time measured until convergence (120 epochs) an NVIDIA GeForce RTX 4090 GPU. The reported time for RVDT additionally includes the training phase of the expert policy Te. As shown in the table, the expert policy training phase typically accounts for approximately one-fifth of the total RVDT training time. When this phase is excluded, the additional overhead introduced by RVDT relative QT is less than 10%. It also be observed that QT itself incurs a substantial increase training time compared to DT, which primarily stems from its additional Q value learning process.
Table 17: Training time comparison measured until convergence. The reported training time for RVDT additionally includes the expert policy training phase.
Gym Tasks
DT
QT
RVDT (total) Expert policy RVDT
halfcheetah-me ---
7.3hr
12.2hr
15.5hr
3.3hr
hopper-m-c
3.6hr
10.3hr
14.6hr
3.0hr
walker2d-m-e
6.1hr
11.7hr
15.4hr
3.3hr
halfcheetah-m
6.3hr
11.8hr
16.1hr hr
3.2hr
hopper-m
3.5hr
9.2hr
14.0hr
3.2hr
walker2d-m
6.1hr
11.5hr
16.0hr
3.2hr
halfcheetah-m-r
7.2hr
12.2hr
17.4hr
3.3hr
hopper-m-r
3.5hr
8.5hr
12.3hr
3.2hr
walker2d-m-r
6.2hr
11.2hr
14.8hr
2.8hr
Average
5.5hr
11.0hr
15.2hr
3.2hr
xxxvii

---

G Discussion
G.1
Imbalance of Datasets
Leanning from imbalanced datasets has long been recognized as a critical challenge in offline re- inforcement learning. Such imbalance typically manifests in datasets predominantly composed of low-return trajectories, with relatively few high-return trajectories. This scenario is common in real-world applications since acquiring high-return trajectories is usually more costly compared to low-return ones.
Existing offiin RL algorithms, including dynamic programming (DP) based methods, imitation learning (IL) based methods, and conditional sequence modeling (CSM) based methods, encounter significant difficulties when learning from imbalanced datasets. DP-based methods often incorporate inductive biases or regularization terms to enforce policy closeness to the behavioral distribution ob- served in the dataset, thereby mitigating out-of--istribun the learned policy to resemble the action distribution of the non-optimal behavior policy. IL-based methods essentially imitate the underlying suboptimal behavior policy that generated the imbal- anced dataset, which can result in overly conservative policies due to excessive imitation of actions from the abundant low-performing trajectories. CSM-based methods typically condition on optimal (or near-optimal) returns during inference for decision-making However, due to the skewed distri- bution of near-optimal versus low-return trajectories in the training datasets, CSM-based methods frequently train trajectories with low returns. This discrepancy in return conditioning between the training and inference phases leads CSM methods to experience severe distributional shift problems when conditioning on high returns for inference.
(OOD) errors [TZ] 9I, which in turn leads
Recent studies describe dataset imbalance from multiple perspectives. For instance, Hong et al. [25] define dataset imbalance through the positive-sided variance of returns (RPSV, defined in Defini- tion MI.. Other research interprets dataset imbalance from the perspective of data coverage, where data coverage broadly encompasses both the coverage of state and action spaces [68] and the trajec- tory return-coverage within the return space [[]].
Definition 1. RPSV of a dataset, V V+lg(T)), corresponds to the second-order moment of the pos- itive component of the difference between trajectory return: g(T:) = Ei- rr((st at) and its expectation, where Ti denotes a trajectory in the dataset:
++1+(Ta)]= = EY [[o(T)) - -E,))((())) with x4 = max(2, xx 0}.
G.2 Expert-Level and Full-Spectrum Return-Coverage
In this paper. we uncover the relationship between the return distribution of the training dataset and the inference- ditrrence-time performance of CSM policies. Since the return distribution of an offline RL dataset corresponds a discrete distribution over the return space, their probability density functions (PDFs) or cumulative distribution functions (CDFs) are challenging to represent. Therefore, this pa- per interprets return distributions from the perspective of return-cvverage. offering a mathematically tractable description suitable for our theoretical analysis.
It commonly understood that for a well-trained CSM algorithm, its inference-time performance largely aligns with the initial input return-togo (RTG) token, which is a manually specified mapping from the initial state 81 real number, i.e., fx1) S IR. The initial RTG token represents the human-expected final return for the test algorithm; thus, it is generally chosen to correspond to the optimal return associated with the optimal policy x² of the underlying MDP. We denote this optimal policy-derived conditioning function as the optimal conditioning function f*. The inference-time returns associated with f" are referred to as the optimal conditional returns. For CSM algorithms, apart from the initial RTG token derived directly from f", subsequent RTG tokens are computed by subtracting the sampled reward from the previous RTG token, i.e., f = f(+) - Due to the inherent randomness of r and the stochastic nature of environmental transitions, subsequent RTG tokens cannot consistently remain optimal [S. IK) Consequently, returns encountered during policy execution follow more general conditioning function f. which with high probability differs from f We term this broader class of conditioning functions as runtime conditioning functions, and the corresponding inference-time returns runtime conditional returns.
xxxviii

---

Based on the preceding descriptions of optimal conditional returns and runtime conditional returns, we define their corresponding coverage as expert-level return-coverage and full-spectrum return- coverage, respectively (consistent with definitions presented in Section I.l.). Let PTy denote the distribution over states, actions, and returns induced by the behavior poliey a, We define the expert- level return-coverage in offline dataset generated by behavior policy Tg as Pa (g (*1))))))) where f corresponds to the optimal conditioning function. Similarly, the full-spectrum return- coverage is defined Pa (g- [(s))(s)), where represents the runtime conditioning function.
Based on these definitions, becomes evident that coveracoverage captures the representation of inference- time conditional returns within the dataset. The central theoretical analysis of this paper (detailed in Section !) aims to explain how expert-level return-coveran covercee and full-spectrum return- coverage collectively influence the performance of CSM policies. Motivated by this theoretical insight, we propose a simple yet effective return-ooverae rebalancing mechanism, designed to sub- stantially enhance the performance of ttate-of-the-art CSM algorithms.
G.3 Limitations and Future Directions
Several aspects merit further exploration. First, RVDT involves training an additional imitation pol- icy for the return-coverage rebalancing process, which introduces modest computational overhead. This cost, however, remains manageable and is well justified by the resulting improvements pol- icy quality. Future work may explore more efficient architectures or parameter-sharinn strategies to further reduce this overhead. Second, the current design of RVDT focuses on the purely offline RL setting. Extending RVDT to online or hybrid settings, where limited online interactions are available, could enable the model to dynamically expand return support beyond the static dataset. In addition, integrating RVDT with data augmentation generative modeling techniques, such as diffusion-basee trajectory stitching or sub-trajectory generation [[2, ( may further enhance re- turn support and improve algorithmic robustness. Third, the current rebalancing mechanism relies on trajectory returns as a measure of trajectory quality. In certain real-world domains, such as health- care or autonomous driving, reliable return estimation can be challenging. In such cases, alternative evaluation criteria or task-specific metries could be incorporated to approximate trajectory quality, thereby broadening the applicability of the proposed framework.
XXXix


Reviews:
[{"review_id": "aGXmSxxwfY", "reviewer": "Reviewer_Lvw8", "summary": "This paper investigates the distributional shift issue in conditional sequence modeling (CSM) for offline RL, particularly caused by imbalanced return distributions. The authors highlight the importance of covering both expert-level and full-spectrum returns, proposing a straightforward return-coverage rebalancing mechanism named Rebalancing Value-regularized Decision Transformer (RVDT). Extensive experiments on D4RL benchmarks validate that RVDT achieves state-of-the-art performance.", "questions": "See the weaknesses part above.", "limitations": "Yes.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n1. The findings presented in this paper, particularly regarding the significant impact of expert-level coverage and full-spectrum conditional returns, are important and valuable to the research community.\n2. Based on the analysis provided above, the proposed RVDT method is both simple and effective.\n3. The experiments conducted are thorough and comprehensive.\n\n**Weaknesses:**\n\n1. Please correct me if I have overlooked anything, but an important ablation study seems to be missing. Specifically, an ablation on the two newly introduced terms in Equation (7)—the policy regularization term and the value-based guidance term—would be helpful. What would be the impact on performance if either of these two terms were removed?\n2. In the experimental section, the authors attribute RVDT's inferior performance compared to D-QL on the larger AntMaze tasks to D-QL's highly expressive policies. However, due to the lack of further experimental evidence, it would be helpful if the authors could clarify how D-QL performs relative to RVDT on larger Maze2D tasks, which are similarly large maze-like environments. Alternatively, could the authors provide additional qualitative or quantitative results to support this claim?\n3. What is the computational overhead introduced by RVDT? For instance, RVDT requires first training an expert policy \\$\\pi\\_e\\$ and then introduces two additional terms during offline RL training.\n4. Some related works also utilize a small portion of expert data to improve offline RL performance \\[1,2,3].\n\n\\[1] Kumar, Aviral, et al. \"Pre-training for robots: Offline RL enables learning new tasks from a handful of trials.\" *arXiv preprint arXiv:2210.05178* (2022).\n\n\\[2] Yang, Qisen, et al. \"Hundreds guide millions: Adaptive offline reinforcement learning with expert guidance.\" *IEEE Transactions on Neural Networks and Learning Systems* 35.11 (2023): 16288–16300.\n\n\\[3] Cheng, Peng, et al. \"Pushing the Limit of Sample-Efficient Offline Reinforcement Learning.\" *ICLR 2025 Workshop on World Models: Understanding, Modelling and Scaling.*", "was_revised": true, "final_justification": "Thank you for your detailed responses, which have addressed most of my concerns. However, the additional computational overhead appears to offer limited cost-effectiveness. I will keep my original rating but increase my confidence score to reflect a leaning toward acceptance.", "rebuttal": "## Response to Reviewer [Lvw8]\n\nWe sincerely thank the reviewer for the constructive and insightful feedback. Below, we address each concern point-by-point.\n\n---\n\n### **Weaknesses 1**\n\n> Please correct me if I have overlooked anything, but an important ablation study seems to be missing. Specifically, an ablation on the two newly introduced terms in Equation (7)—the policy regularization term and the value-based guidance term—would be helpful. What would be the impact on performance if either of these two terms were removed?\n\n**Response:**  \nWe are happy to provide clarification on this point. In fact, the ablation study you mentioned is included as part of our analysis in Section 5.3. The two newly introduced terms, (1) the policy regularization term and (2) the value-based guidance term, correspond to the Explicit Rebal. (Explicit Rebalancing Mechanism) and Implicit Rebal. (Implicit Rebalancing Mechanism) presented in Table 5, respectively. In Table 5, the variant “RDT” represents the ablated version of RVDT without the value-based guidance component, while “VDT” corresponds to the version without the policy regularization term.\n\nTable 5: Component-level breakdown across ablation variants\n|Component|DT|DT-Dup|**RDT**|QT|**VDT**|RVDT-Dup|RVDT-Determ|RVDT|\n|-|-|-|-|-|-|-|-|-|\n|Explicit Rebal.|None|Dup.|**KL**|None|**None**|Dup.|KL|KL|\n|Implicit Rebal.|None|None|**None**|Q-value|**Q-value**|Q-value|Q-value|Q-value|\n|Policy Type|Determ.|Stoch.|**Stoch.**|Determ.|**Stoch.**|Stoch.|Determ.|Stoch.|\n\nThe impact of removing either component can be observed in Table 6 and is discussed in Lines 344–347. Specifically, as stated: “Both RDT (explicit rebalancing) and QT/VDT (implicit rebalancing) individually outperform DT, confirming the effectiveness of applying return-coverage rebalancing mechanisms on DT, and indicating that the combination of rebalancing mechanisms in RVDT provides additional performance gains.”\n\nTable 6: Performance comparison of ablation variants across MuJoCo tasks\n|Task|DT|DT-Dup|**RDT**|QT|**VDT**|RVDT-Dup|RVDT-Determ|RVDT|\n|-|-|-|-|-|-|-|-|-|\n|halfcheetah|84.2|90.3|**90.5**|91.2|**89.5**|93.4|91.5|94.9|\n|hopper|109.5|112.1|**111.9**|112.3|**112.6**|112.1|113.6|113.8|\n|walker2d|108.2|108.8|**109.7**|113.2|**110.3**|110.9|113.1|118.7|\n|Average|100.6|103.7|**104.0**|105.6|**104.1**|105.5|106.1|109.1|\n\nWe will revise the manuscript to make this correspondence clearer.\n\n---\n\n### **Weaknesses 2: RVDT vs. D-QL in Larger Maze-like Environments**\n\n> In the experimental section, the authors attribute RVDT's inferior performance compared to D-QL on the larger AntMaze tasks to D-QL's highly expressive policies. However, due to the lack of further experimental evidence, it would be helpful if the authors could clarify how D-QL performs relative to RVDT on larger Maze2D tasks, which are similarly large maze-like environments. Alternatively, could the authors provide additional qualitative or quantitative results to support this claim?\n\n**Response:**  \nWe sincerely thank the reviewer for the thoughtful observation and helpful suggestion. \nActually, compared to RVDT, D-QL leverages a more expressive diffusion policy, but lacks the long-term modeling capability offered by conditional sequence modeling and the return-coverage rebalancing mechanism we propose.\nTo provide additional clarity, we present the performance of D-QL on larger Maze2D tasks below:\n\nTable: Performance comparison on Maze2D Tasks.\n|Maze2D Tasks|CQL|IQL|D-QL|RVDT|\n|-|-|-|-|-|\n|maze2d-m|41.8|34.9|91.0|183.5|\n|maze2d-l|49.6|61.7|200.7|254.3|\n\nFrom the results above, it can be clearly observed that the D-QL algorithm, which utilizes a diffusion policy, significantly outperforms other Q-learning-based algorithms with simpler Gaussian policies, i.e., CQL and IQL. This highlights the superiority of diffusion policies over Gaussian policies.\n\nNote that, unlike in the AntMaze task, RVDT outperforms D-QL on the Maze2D tasks. This discrepancy may be attributed to the less complex agent used in Maze2D, which is a simpler particle with a relatively small action space. In contrast, the multi-legged agent in AntMaze places greater demands on multi-modality action modeling capabilities. As a result, in the Maze2D tasks, the performance gains from the diffusion policy are less pronounced than those provided by long-term modeling and the return-coverage rebalancing mechanism.\n\nBased on these insights, we conclude that both highly expressive policy architectures and our proposed return-coverage rebalancing mechanism contribute positively to overall algorithm performance. Particularly, in tasks characterized by multi-modal action spaces (e.g., AntMaze), the expressiveness of the policy becomes especially crucial. This conclusion aligns well with the prevailing research direction that emphasizes policy expressiveness. We leave further exploration of combining diffusion policies with return-coverage rebalancing in CSM frameworks as our future work. \n\n---\n\n### **Weaknesses 3: Computational Overhead of RVDT**\n\n> What is the computational overhead introduced by RVDT? For instance, RVDT requires first training an expert policy $\\pi_e$ and then introduces two additional terms during offline RL training.\n\n**Response:**  \n\nWe sincerely thank the reviewer for this careful observation. Below, we provide a table detailing the training time comparison of DT, QT, RVDT (including the expert policy $\\pi_e$ training phase), and the expert policy training phase itself, measured until the algorithm has converged (i.e., 120 training epochs) on an NVIDIA GeForce RTX 4090 GPU.\n\nTable: Training time comparison, where RVDT's training time includes the expert policy training phase.\n|Gym Tasks|DT|QT|RVDT (total)|Expert policy|\n|-|-|-|-|-|\n|halfcheetah-m-e|7.3hr|12.2hr|15.5hr|3.3hr|\n|hopper-m-e|3.6hr|10.3hr|14.6hr|3.0hr|\n|walker2d-m-e|6.1hr|11.7hr|15.4hr|3.3hr|\n|halfcheetah-m|6.3hr|11.8hr|16.1hr|3.2hr|\n|hopper-m|3.5hr|9.2hr|14.0hr|3.2hr|\n|walker2d-m|6.1hr|11.5hr|16.0hr|3.2hr|\n|halfcheetah-m-r|7.2hr|12.2hr|17.4hr|3.3hr|\n|hopper-m-r|3.5hr|8.5hr|12.3hr|3.2hr|\n|walker2d-m-r|6.2hr|11.2hr|14.8hr|2.8hr|\n|**Average**|**5.5hr**|**11.0hr**|**15.2hr**|**3.2hr**|\n\nAs illustrated in the table, the expert policy training phase typically accounts for roughly one-fifth of the total RVDT training time. Excluding this phase, the remaining training overhead introduced by RVDT relative to QT is less than 10%. In addition, we observe that QT already introduces a significant increase in training time compared to DT.\n\nIt is important to note that the primary objective of our work is to derive high-performing policies from fixed datasets. As a result, the extended training time incurred by RVDT represents a modest trade-off, especially given the significant performance gains it delivers over other CSM-based methods. The resulting improvements in policy quality outweigh the additional computational cost.\n\n---\n\n### **Weaknesses 4: Related Work on Expert Data Usage**\n\n> Some related works also utilize a small portion of expert data to improve offline RL performance [1,2,3].\n\n**Response:**  \n\nWe thank the reviewer for bringing these relevant references to our attention. These works expand the scope of expert data usage in offline RL across different settings.\n\nSpecifically, **Kumar et al.** propose PTR [1], a meta-learning-inspired framework that first performs multi-task offline RL pretraining on diverse robotic demonstration datasets, then fine-tunes on target tasks using only 10-15 expert demonstrations. **Yang et al.** propose GORL [2], a plug-in framework that adaptively balances policy improvement and constraint for each training sample by employing a guiding network trained with only a small number of expert demonstrations. The guiding network is meta-optimized in a MAML-like manner, alternating between updates on expert data and policy learning on the offline dataset. **Cheng et al.** propose TELS [3], which aims to enable effective policy learning in low-data regimes (e.g., 1% of the original dataset). In summary, these settings include pretraining and fine-tuning approaches for robotics [1], meta-learning approaches [1,2], and sample-efficient offline RL with limited data [3].\n\nOur method focuses on leveraging expert data in a more general offline RL setting, which differs from the settings for these works. We acknowledge that these works also explore the utilization of expert data under more complex or specialized scenarios. We will add an additional discussion of methods that utilize a small portion of expert data under broader or task-specific problem formulations in our revision.\n\n[1] Kumar, Aviral, et al. \"Pre-training for robots: Offline RL enables learning new tasks from a handful of trials.\" arXiv preprint arXiv:2210.05178 (2022).\n\n[2] Yang, Qisen, et al. \"Hundreds guide millions: Adaptive offline reinforcement learning with expert guidance.\" IEEE Transactions on Neural Networks and Learning Systems 35.11 (2023): 16288–16300.\n\n[3] Cheng, Peng, et al. \"Pushing the Limit of Sample-Efficient Offline Reinforcement Learning.\" ICLR 2025 Workshop on World Models: Understanding, Modelling and Scaling.\n\n---\n\nAgain, we greatly appreciate the reviewer’s helpful suggestions and the opportunity to clarify and improve our work. We hope our response can alleviate some of your concerns and offer useful clarification."}, {"review_id": "qsl4JMY7P6", "reviewer": "Reviewer_8NTa", "summary": "This paper explore the CSM used in offline RL. The auther studies two types of return (expert-level and full spectrum) and provide some theoretical analysis on their impact. Additionally, auther provide a return converage rebalancing mechanism into DT framework.", "questions": "1. Can auther discuss the validity of the return-coverage assumption? Those assumptions may hold in small or well-curated offline datasets, it becomes increasingly difficult to satisfy as the dataset grows in size and complexity. In such cases, $\\alpha$ may become vanishingly small, which would cause the theoretical performance bounds to become vacuous or uninformative. \n\n2. How the quality of the expert policy affects RVDT performance, How long it takes to train the expert policy, and the trade-off between expert policy quality, training time, and RVDT effectiveness?\n\n3. For Table 1 gym tasls, majority of RVDT performance are so close to the QT  if we consider QT's statistic value, it may hard to tell if RVDT has performance different with QT for the gym tasks. Same issue for Kitchen tasks and AntMaze tasks\n\n4. Why is  top 50% of trajectories to construct the expert dataset not more?\n\nI would be willing to raise my rating if the authors can address 1. the assumption bound issue 2. the inconsistance statistic result report, as they are my primary concern and the main reason for my initial assessment.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": ">Strength\n\n1. the paper is clear\n2. auther did compare with various benchmarks and algorithms\n\n>Weaknesses\n\n1. The major concern is the assumptions in the theorem and the bound value $\\alpha_f$. In order for those assumption to hold, the $\\alpha$ maybe really small or even close to 0.  If this is the case, the bound for theorem 1 and lemma 1 become meaningless ($\\infty$).\n\n2. The method relies on training an expert policy via imitation learning on near-optimal trajectories, which is then used to regularize the main policy through a KL divergence term. The quality of the expert policy will affect the performance of RVDT. Additionally, the higher quality of expert policy will result in longer pretraining time and utilize larger dataset.\n\n3. For the evaluation, some of results have statistic value (std) some are not.", "was_revised": true, "final_justification": "The reviewer addressed all concerns; however, I cannot give a full acceptance score (5) because relying on an assumption or theorem with clear limitations—regardless of its common use in other works—is not a strong basis for theoretical analysis.", "rebuttal": "## Response to Reviewer [8NTa]\nWe sincerely thank the reviewer for the thoughtful and constructive feedback. We address each point below in detail.\n\n---\n\n### **Weakness 1 and Question 1**\n**Response:**  \nWe provide clarifications to address the reviewer’s concern regarding the return-coverage assumption:\n\n1. **Alignment with existing literature.**\nThe return-coverage assumption we adopt is a standard analytical assumption in most existing CSM-based methods, including works such as DT [1] and QT. In fact, theoretical analyses in most existing offline RL algorithms typically rely on some form of coverage condition, either over return pairs or the state-action space, to ensure that the learning problem is well-posed and admits meaningful generalization guarantees.\n2. **Interpretation of the theoretical bound.**\nWe fully acknowledge the reviewer’s valid concern regarding extremely small return-coverage. In such cases, the theoretical performance gap becomes unbounded, which actually reflects the inherent nature of CSM-based methods. Specifically, when the return-coverage of the training data is insufficient, reliable policy learning cannot be guaranteed, and the performance gap can become arbitrarily large. This phenomenon, as predicted by our theory, has already been widely observed in empirical studies: the more comprehensively a dataset covers expert-level and full-spectrum conditional returns, the better the performance achieved by CSM-based methods.\n3. **The necessity of coverage in offline RL.**\nNote that the return-coverage required in our analysis is not strictly restricted to optimal trajectories. Indeed, coverage over relatively high-quality or expert-level returns is often sufficient to ensure a meaningful performance gap. Our analysis naturally extends to datasets containing only a subset of high-quality trajectories, which can guide rebalancing in CSM-based methods. The necessity of coverage over relatively high-quality trajectories aligns with common intuition in offline RL. For example, for BC-based methods, generalization depends on sufficient coverage over useful state-action pairs. Similarly, CSM-based methods require non-negligible return-coverage to achieve satisfactory performance. In contrast, if the dataset contains only uniformly poor trajectories, no CSM-based method can be expected to recover a high-quality policy.\n4. **Practical validity.**\nIn practice, as the dataset grows in size, the factors that primarily influence return-coverage are the difficulty of the task and the data collection strategy. Generally, a good data collection strategy is required to obtain sufficient return-coverage for difficult tasks. Actually, widely adopted large-scale offline RL benchmarks (e.g., D4RL) typically include a non-trivial amount of high-quality data, which satisfy the coverage assumption. Practically, high-quality trajectories may be diluted by a large quantity of lower-quality data as the dataset grows. Our rebalancing mechanism plays a crucial role in dynamically adjusting return coverage, thereby keeping the performance gap within a relatively low range.\n\n[1] Brandfonbrener, David, et al. When does RCSL work for offline rl. NIPS. 2022.\n\n---\n\n### **Weakness 2 and Question 2**\n**Response:**  \nIndeed, the quality of the expert policy influences the performance of RVDT. Nevertheless, we emphasize that the expert policy is trained using a subset $D_e \\subset D$ rather than any additional data. It is worth noting that the \"expert data\" in our method does not necessarily refer to exact expert demonstrations, but rather to relatively high-quality trajectories selected from the existing dataset.\n\nWe empirically investigate how expert policy quality impacts RVDT performance. As shown in Table 2, adding more expert trajectories to a fixed non-expert dataset (RVDT(0%)–RVDT(50%)) consistently improves performance. A higher percentage indicates more expert data and thus a higher-quality expert policy used for regularization. See Section 5.2 (lines 301–310) for details.\n\nTable 2: Performance with increasing expert trajectory proportions.\n|Task|RVDT(0%)|RVDT(15%)|RVDT(30%)|RVDT(50%)|\n|-|-|-|-|-|\n|halfcheetah|49.3|92.7|93.3|94.4|\n|hopper|100.2|113.1|113.8|113.9|\n|walker2d|92.7|110.3|112.4|113.8|\n|Average|80.3|105.4|106.5|107.0|\n\nTable 2 shows that RVDT performance consistently improves with better expert policy quality, with the largest gain observed between RVDT(0%) and RVDT(15%). The results reveal the relationship between expert policy quality and RVDT performance, demonstrating both the effectiveness and benefit of our method.\n\nRegarding **training time**, the dominant influencing factor is the number of training epochs, rather than the dataset size. The training time is reported in the table below. The expert policy training time for RVDT was approximately 3 hours on an NVIDIA 4090 GPU, with minimal variation across different tasks.\n\nTable: Training time comparison, where RVDT's training time includes the expert policy training phase.\n|Gym Tasks|DT|QT|RVDT (total)|Expert policy|\n|-|-|-|-|-|\n|halfcheetah-m-e|7.3hr|12.2hr|15.5hr|3.3hr|\n|hopper-m-e|3.6hr|10.3hr|14.6hr|3.0hr|\n|walker2d-m-e|6.1hr|11.7hr|15.4hr|3.3hr|\n\nTo address the **trade-off** between expert policy quality, training time, and RVDT effectiveness, we primarily rely on tuning several key hyperparameters within the algorithm.\nNote that our primary goal is to find a policy that achieves the highest possible performance given fixed training data. Therefore, rather than explicitly trading off between the three, we search for optimal hyperparameters. Specifically, RVDT selects the top $\\rho$% of high-quality trajectories to train the expert policy until convergence. Then, it adjusts the coefficients $\\eta$ and $\\alpha$ in the overall loss function: $L(RVDT) = L(DT) - \\eta L(Q) + \\alpha L(KL)$. In practice, we perform grid search over $\\rho$, $\\eta$, and $\\alpha$, and select the combination that yields the best performance. \n\n---\n\n### **Weakness 3 and Question 3**\n**Response:**  \nWe appreciate the reviewer’s emphasis on statistical rigor. Following common convention in CSM-based offline RL research (e.g., DT, CQL, QT), we report both the mean and standard deviation (std.) for our own method, while only reporting mean values for other baselines, consistent with prior works. Note that some previous works do not report std., and it is generally infeasible for us to report std. for all baselines.\n\nNevertheless, given the importance of QT, we will include its std. in the revised manuscript. A subset of the results is shown in the table below due to space limitations. The results demonstrate that RVDT consistently outperforms QT, even when accounting for the variance. Notably, QT exhibits higher variance than RVDT, largely due to its sensitivity to Q-value estimation errors. In contrast, the KL regularization term in RVDT helps mitigate the impact of such errors. Furthermore, according to Proposition 2 (lines 241–244), the performance of RVDT would generally be better than that of QT.\n\n|Gym Tasks|QT|RVDT|\n|-|-|-|\n|halfcheetah-m-e|93.2 ± 0.8|94.4 ± 0.1|\n|hopper-m-e|113.0 ± 0.2|113.1 ± 0.5|\n|walker2d-m-e|112.0 ± 0.3|112.7 ± 1.6|\n|...|...|...|\n|pen-human|111.9 ± 11.8|127.2 ± 5.5|\n|hammer-human|10.4 ± 5.3|24.0 ± 1.5|\n|pen-cloned|85.8 ± 12.6|117.8 ± 8.6|\n|hammer-cloned|11.8 ± 5.8|21.3 ± 2.7|\n|...|...|...|\n|antmaze-u-d|92.0 ± 6.2|98.0 ± 4|\n|antmaze-m-d|24.0 ± 8.6|30.0 ± 6.3|\n|antmaze-l-d|10.0 ± 0|10.0 ± 0|\n\nRegarding the reviewer’s concern that RVDT shows only modest gains over QT on certain tasks (e.g., MuJoCo), this is mainly due to boundary effects: both methods already reach or exceed the dataset’s expert-level returns, leaving any further improvement inherently challenging.\n\nRVDT’s advantage over QT is most evident when accurate Q-value estimation is difficult, such as in low-data or sparse-reward settings. As shown in Table 3 (lines 311–320), RVDT significantly outperforms QT under low-data regimes. The four datasets used here contain only about one-tenth of the original D4RL data, with difficulty increasing from $D_1$ to $D_4$ (see Appendix F.1 for details). RVDT consistently outperforms QT, achieving average improvements of 45.4%, 14.5%, 14.2%, and 45.6% over QT on $D_1$–$D_4$, respectively. Additional evidence of RVDT’s advantage in sparse-reward settings is shown in Table 4 (lines 321–329), where it outperforms QT by 27%, 38%, 7%, and 6% on maze2d-open, maze2d-umaze, maze2d-medium, and maze2d-large, respectively. Collectively, these results confirm RVDT’s clear and consistent superiority over QT under challenging conditions.\n\nTable 3: Performance comparison in low-data regimes.\n|Task|QT($D_1$)|RVDT($D_1$)|QT($D_2$)|RVDT($D_2$)|QT($D_3$)|RVDT($D_3$)|QT($D_4$)|RVDT($D_4$)|\n|-|-|-|-|-|-|-|-|-|\n|maze2d-u|100|171|81|101|73|77|61|100|\n|maze2d-m|137|187|175|190|163|182|98|175|\n|maze2d-l|109|140|81|90|104|131|100|95|\n|Average|115|166|112|127|113|130|86|124|\n\n---\n\n### **Question 4**\n\n**Response:**  \nIn fact, the percentage of top $\\rho$% trajectories selected to construct the expert dataset is a tunable hyperparameter. For our main results reported in Table 1 (Section 5.1 of the manuscript), the specific values used for this hyperparameter are provided in Appendix E (Table 10), and they vary across tasks. This hyperparameter is selected via grid search.\n\nFor ablation studies, unless otherwise specified, we adopt the 50% setting primarily for the purpose of controlled comparison. One reason for choosing 50% is that it represents the most commonly used value across tasks (see Table 10), making it a representative and balanced choice for the general analysis.\n\n---\n\nAgain, we thank the reviewer for their valuable comments and suggestions. We hope our response can alleviate some of your concerns and offer useful clarification."}, {"review_id": "hot6wyzSsh", "reviewer": "Reviewer_4BE1", "summary": "This paper addresses the trade-off between return maximization and behavior coverage in conditional sequence modeling (CSM)-based offline RL methods such as Decision Transformer. They propose a test-time control framework that balances this trade-off by reranking sampled trajectories using a weighted combination of predicted return and a learned coverage score. Their method, which does not alter the training process, shows improved performance across several D4RL tasks.", "questions": "1. The combined use of Q-value regularization and expert KL divergence may result in conflicting optimization signals, particularly when Q-value estimation is poor or inaccurate. How do you address or mitigate the potential instability and performance degradation that may arise from this issue?\n2. Besides using top-k return for expert trajectory selection, did the authors consider or evaluate any alternative definitions or methods for identifying expertise in the offline dataset?\n3. The bound in Theorem 2 shows a dependence on both expert-level and full-spectrum return coverage parameters (α_f and α_f*). Could you provide more intuition or empirical evidence about how trade-offs between these two types of coverage affect policy performance? Are there scenarios where improving one necessarily degrades the other?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper addresses a meaningful and underexplored problem in conditional sequence modeling (CSM) for offline RL—namely, the trade-off between return maximization and behavioral coverage—which is both practically relevant and conceptually important.\n2. The theoretical formulation is relatively complete and provides a useful lens to understand the inherent tension in CSM-based policies between imitating high-return behaviors and maintaining adequate state-action coverage.\n3. The experimental design is comprehensive, evaluating across multiple dimensions including data quality (expert vs mixed), data scale (low-data regime), and reward density (sparse vs dense).\n\nWeaknesses:\n\n1. The proposed rebalancing mechanism fundamentally relies on the availability and selection of high-quality, expert-level trajectories. In many real-world offline RL settings, such as medical or autonomous driving applications, “expert” behavior is often ambiguous.\n2. While the method claims to enhance expert-level coverage, in practice the KL regularization simply increases the weight of expert-like behaviors, but does not expand the policy’s support into previously unobserved state-action regions.", "was_revised": false, "final_justification": "", "rebuttal": "## Response to Reviewer [4BE1]\n\nWe sincerely thank the reviewer for the detailed and insightful comments. We address each question below.\n\n---\n\n### **Weakness 1**\n\n> The proposed rebalancing mechanism fundamentally relies on the availability and selection of high-quality, expert-level trajectories...\n\n**Response:**  \n\nWe sincerely thank the reviewer for raising this important point. Note that the return-coverage required in our analysis is not strictly restricted to optimal trajectories. In practice, the presence of relatively high-quality trajectories or near-expert returns is often sufficient to facilitate meaningful learning outcomes. We would like to clarify that our method does not require any externally provided expert data, and it operates entirely on the given offline dataset and remains effective as long as there is sufficient diversity to enable relative ranking. As demonstrated in Table 1 of the main paper, our method consistently improves performance even on the \"-medium\" and \"-medium-replay\" datasets, which contain no expert-level trajectories, thereby highlighting its practicality beyond idealized settings.\n\nNevertheless, we agree that in certain real-world domains, such as healthcare or autonomous driving, assessing trajectory quality based solely on return may be challenging. In such cases, additional annotations or task-specific metrics may be necessary to approximate trajectory quality, which could then serve as a foundation for applying our rebalancing framework. We will include a discussion of these potential extensions in the revision and consider this an important direction for future work.\n\n---\n\n### **Weakness 2**\n\n> The KL regularization simply increases the weight of expert-like behaviors, but does not expand the policy’s support...\n\n**Response:**  \n\nWe appreciate the reviewer’s insightful observation. Generally speaking, there are two aspects that influence the return-coverage: reward coverage, which is explicitly involved in the definition, and state-action space coverage, which is implicitly involved. To increase return-coverage, we introduce an expert KL regularization that effectively performs weighted sampling over a proportion of high-quality trajectories with high cumulative rewards. As for state-action space coverage, it is generally difficult to expand the policy’s support into previously unobserved state-action regions in offline RL without resorting to data augmentation or generative modeling. Indeed, prior work [1,2] has explored this direction by leveraging model-based approaches or generative models (e.g., diffusion models) to synthesize new transitions via trajectory stitching or sub-trajectory generation. These approaches are discussed in Lines 97–101.\n\nIt is important to emphasize that our proposed rebalancing mechanism is orthogonal and complementary to these methods. It can be seamlessly applied on top of augmented or generated datasets to further expand policy support while still leveraging expert KL regularizer. We view the integration of rebalancing mechanisms with advanced data generation techniques as a promising direction for future work.\n\n[1] Guanghe Li, et al. Diffstitch: Boosting offline rl with diffusion-based trajectory stitching. ICML, 2024.\n\n[2] Sungyoon Kim, et al. Stitching sub-trajectories with conditional diffusion model for goal-conditioned offline rl. AAAI, 2024.\n\n---\n\n### **Question 1**\n\n> The combined use of Q-value regularization and expert KL divergence may result in conflicting optimization signals...\n\n**Response:**  \nWe sincerely thank the reviewer for this insightful observation. As described in Eq. (7) of the paper, the loss function of RVDT can be summarized as:\n$$L(RVDT) = L(DT) - \\eta L(Q) + \\alpha L(KL).$$\n\n**From an empirical perspective**, we mitigate potential instability by carefully balancing the hyperparameters $\\eta$ and $\\alpha$ associated with the two regularizers. In practice, for datasets where the Q-values are difficult to estimate accurately, such as those with sparse rewards, we typically set $\\eta$ to a relatively small value. This shifts the optimization emphasis toward the expert KL regularizer, which serves as a more reliable guiding signal and helps stabilize training in the presence of Q-value estimation errors. \n\n**From a theoretical perspective**, Proposition 1 in our paper provides partial justification for the compatibility of the two regularizers. Specifically, under the Gaussian policy assumption, the expert KL regularizer can be interpreted as implementing a weighted sampling scheme, leading to the following reformulated loss:\n$$ L_{RVDT}(θ) =  E_{τ \\sim D} [ (1 + α · I[τ ∈ D_e]) · ( Σ_{i=1}^H -log π_θ(a_i | s_i, g(τ_i), τ̄_{t-1}^K) ) ] - \\eta L(Q).$$\nThis formulation suggests that the expert KL regularizer effectively operates through a sampling-based mechanism that is inherently compatible with the main DT loss. As such, it does not introduce significant conflict during optimization and maintains stability in practice.\n\nRegarding the concern of **performance degradation**, the dual-regularizer design of RVDT is specifically intended to enhance robustness in challenging settings, such as those with sparse or unreliable reward signals where Q-value estimation is inherently difficult. Our experimental results demonstrate that RVDT consistently outperforms QT in such scenarios, validating the practical effectiveness of this design. Furthermore, as shown in Proposition 2 (lines 241–244), the performance of RVDT is generally higher than that of QT.\n\nIn summary, by appropriately tuning $\\eta$ and $\\alpha$, RVDT effectively mitigates potential conflicts between regularizers and avoids the instability or performance degradation.\n\n---\n\n### **Question 2**\n\n> Did the authors consider or evaluate any alternative definitions or methods for identifying expertise in the offline dataset?\n\n**Response:**  \n\nWe sincerely thank the reviewer for this thoughtful question. Actually, alternative criteria for assessing trajectory quality, such as the averaged Q-value along the trajectory and Conditional Value-at-Risk, can also be adopted. However, these alternatives are more complicated to compute and typically require additional neural network approximation, thus making them less well suited to our current setting and tasks than trajectory return.\n\nOur study primarily focuses on the theoretical analysis of how the performance of conditional sequence modeling (CSM)-based policies is affected by the return-coverage of expert-level and full-spectrum trajectories in the offline training dataset. Building upon our theoretical results, we propose a simple yet effective return-coverage rebalancing mechanism to improve the performance of current CSM-based algorithms, and this mechanism also serves as a plug-in module that can be easily integrated into existing CSM-based algorithms. In this context, using trajectory return as a proxy for expertise is a straightforward criterion within the standard setting of CSM-based algorithms. It aligns well with the return-conditioned training paradigm and provides a natural way to stratify trajectories by quality.\n\nNevertheless, exploring alternative expert selection criteria, especially in more complex scenarios such as reward-free environments and constrained MDPs, would indeed be meaningful. However, such settings are beyond the scope of the present work. We consider this a promising direction and leave it to future research.\n\n---\n\n### **Question 3**\n\n> Could you provide more intuition or empirical evidence about how trade-offs between these two types of coverage affect policy performance?... \n\n**Response:**  \n\nWe sincerely thank the reviewer for the detailed observation. The performance bound of a CSM-based policy is jointly influenced by $\\alpha_f$ and $\\alpha_f^{\\star}$. In the offline RL setting, the relationship $\\alpha_f + \\alpha_f^{\\star} \\leq 1$ holds when $f \\neq f^{\\star}$, implying that increasing one is likely to result in a decrease in the other. To empirically validate this theoretical insight, we present an experiment in Appendix F.2 (lines 1121-1132), with results summarized in Table 12. This experiment evaluates RVDT’s performance when return-coverage rebalancing is applied using varying proportions of top-$k$ trajectories based on return.\n\nTable 12: Performance comparison with different proportions of subsampled expert datasets.\n|Task (sparse R)|10% RVDT|30% RVDT|50% RVDT|70% RVDT|100% RVDT|\n|-|-|-|-|-|-|\n|maze2d-u-v1|93.1|**143.0**|100.4|102.4|101.6|\n|maze2d-m-v1|189.9|**195.0**|192.5|191.6|188.6|\n|maze2d-l-v1|246.5|**250.0**|248.5|248.1|243.6|\n|Average|176.50|**196.0**|180.47|180.70|177.93|\n\nAs shown in the table, varying the proportion of top-$k$ trajectories directly affects the trade-off between $\\alpha_f^{\\star}$ and $\\alpha_f$: as the proportion increases from 10% to 100%, $\\alpha_f^{\\star}$ first increases and then decreases, while $\\alpha_f$ follows the opposite trend. Correspondingly, the performance of RVDT peaks when 30% of the highest-return trajectories are used, suggesting that $\\alpha_f$ and $\\alpha_f^{\\star}$ reach a near-optimal balance under this setting.\n\nThis finding supports our theoretical results and emphasizes that effective return-coverage rebalancing is not achieved by simply maximizing the use of expert trajectories. Intuitively, relying exclusively on expert data (i.e., maximizing $\\alpha_f^{\\star}$ while neglecting $\\alpha_f$) can lead to significant performance degradation due to distributional shift, especially when the training data is limited. Therefore, careful calibration and balancing of both expert-level and full-spectrum return-coverage is essential for achieving optimal performance in offline RL.\n\n---\n\nOnce again, we thank the reviewer for their thoughtful feedback, and we hope our response can alleviate some of your concerns and offer useful clarification."}, {"review_id": "qhXH5Fyits", "reviewer": "Reviewer_Mnay", "summary": "One can roughly categorize offline RL algorithms as either \"reweighted imitation learning\" or \"Q learning constrained to in-distribution data\". The authors focus on the former, in a sub-area called Conditional Sequence Modeling (CSM). In CSM, we condition the trajectory on a desired Reward to Go (RTG) $g$, denoting the remaining reward in a trajectory from time $t$: $\\tau_t$\n\n$$ argmin_\\pi - \\log \\pi(a_t | s_t, g(\\tau_t), \\tau_{t-1}).  $$\n\nIt is difficult to predict $g$. Instead, one can learn what is effectively a value function $f_\\pi(s) = V_\\pi(s)$\n\n\n$$ argmin_\\pi - \\log \\pi(a_t | s_t, f_\\pi(s_t), \\tau_{t-1}).  $$\n\nThus, the problem reduced to finding the actions that lead to a specific return, where the (undiscounted) return is specified by $f$. By setting $f = V_*$, we can recover an optimal policy.\n\nFirst, the authors analyze the performance gap caused by non-optimal $f$, (i.e., $f \\neq V_*$). Then, they provide sample complexity bounds for CSM. Next, they provide the first algorithmic contribution of the paper, termed RDT.\n\nIn RDT, the authors select an expert subset of trajectories and learn a policy on this subset using imitation learning. More formally, they learn a policy that learns a return-conditioned policy on the full dataset, while also imitating the expert-subset policy. In other words, the authors can leverage the full dataset for training, while biasing the policy towards \"better\" trajectories. The authors utilize RDT to propose RVDT, the main contribution of the paper.\n\nRVDT takes the RDT objective and adds a term that represents the Q function. The full RVDT objective contains three terms -- two from RDT and a new Q value term:\n- The return-to-Go objective\n- The imitation learning objective, for the expert subset of trajectories\n- Maximization of $Q_\\pi$\n\nThe authors test their algorithm across standard offline RL benchmarks, demonstrating that their approach is state of the art. They perform ablations demonstrating that their method scales well with increasingly-expert data, as well as in low-data regimes.", "questions": "- Contributions: Each contribution is a bit long and difficult to parse without reading the preliminary section. Could you simplify/summarize these a bit? \n- Section 3.2: Can you please add a sentence with a more intuitive explanation of the conditioning function $f$? The $f$ we aim to find is often the RTG of the optimal policy $\\pi_*$, correct? $f(s) = \\max_\\pi \\mathbb{E}_{\\pi}[g(s)]$?", "limitations": "I do not think the authors list any limitations. I would suggest they consider adding a few sentences about this. One limitation is that their method requires learning an extra imitation learning policy, which can be costly. Especially in scenarios where the performance gap betwee QT and RVDT is negligible.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 2, "strengths_and_weaknesses": "### Strengths\n- The paper is very well written and I enjoyed reading it. It tackles many complex topics, yet the writing style made it easy to follow along.\n- Theorem 2 and 3 provide good theoretical as to why the authors' method should succeed\n- The RVDT objective seems fairly simple to implement, as it is just the combination of a few common objectives\n- RVDT performs very well compared with prior work\n- The authors provide sufficient ablations and analysis of their algorithm\n\n### Weaknesses\n- My biggest concern is that the author's method is very similar to QT. My understanding is that their objective just has one more term, corresponding to the imitation learning objective from an expert subset. Many experiments in Table 1 demonstrate that the performance of QT is only marginally less than RVDT.\n- Tuning hyperparameters $\\eta, \\alpha$ seems like it could be difficult. The automatic $\\eta$ normalization scheme still likely requires hyperparameter tuning based on the variance of Q values.", "was_revised": true, "final_justification": "The authors resolved my main concern about the originality of their approach by pointing out the theoretical contributions. I would improve my originality score to a 3 or 4, but this does not seem possible right now.", "rebuttal": "## Response to Reviewer [Mnay]\n\nWe thank the reviewer for the detailed feedback and thoughtful comments. Below, we address each identified concern point-by-point.\n\n---\n\n### **Weakness 1: RVDT is similar to QT**\n\n> My biggest concern is that the author's method is very similar to QT. My understanding is that their objective just has one more term, corresponding to the imitation learning objective from an expert subset. Many experiments in Table 1 demonstrate that the performance of QT is only marginally less than RVDT.\n\n**Response:**  \n\nWe sincerely thank the reviewer for the detailed feedback. The primary contribution of our work lies in the theoretical investigation of how the performance of CSM-based policies is influenced by the return-coverage of both expert-level and full-spectrum conditional returns in offline reinforcement learning. Based on our theoretical analysis, we propose a simple yet effective rebalancing mechanism. Actually, our theoretical analysis also provides a foundation for other approaches, such as prioritized experience replay, importance reweighting, and advantage-weighted sampling. The expert KL regularization term is introduced as a lightweight, plug-in component that can be seamlessly integrated into existing CSM-based algorithms to enhance robustness and performance, including QT. Moreover, Proposition 2 also shows that the performance upper bound of the optimal policy induced by RVDT’s objective is generally higher than that of QT.\n\nEmpirical results confirm the efficacy of the simple return-coverage rebalancing mechanism: the added KL regularization consistently improves RVDT upon QT, especially in challenging settings where Q-value estimates are unreliable. As provided in Table 1 of our main paper, RVDT outperforms QT by approximately 20% on Maze2D tasks with sparse rewards. In addition, Section 5.2 (Lines 301–329) presents three targeted experiments designed to demonstrate RVDT’s superiority over QT in terms of: (1) effective utilization of limited expert data, (2) performance under low-data regimes, and (3) sparse reward environment. The corresponding results are reported in Tables 2, 3, and 4. We reproduce Table 3 below as an illustrative example to highlight the superiority of RVDT in low-data regimes.\n\nTable 3: Performance comparison in low-data regimes.\n|Task|QT($D_1$)|RVDT($D_1$)|QT($D_2$)|RVDT($D_2$)|QT($D_3$)|RVDT($D_3$)|QT($D_4$)|RVDT($D_4$)|\n|-|-|-|-|-|-|-|-|-|\n|maze2d-umaze|100.3|171.7|81.8|101.8|73.1|76.9|61.4|100.5|\n|maze2d-medium|137.1|187.4|175.2|190.0|163.2|182.0|98.3|175.3|\n|maze2d-large|109.5|140.4|81.5|90.1|104.4|131.3|100.2|95.1|\n|Average|115.6|166.5|112.8|127.3|113.6|130.1|86.6|123.6|\n\n---\n\n### **Weakness 2: Difficulty of Hyperparameter Tuning**\n\n> Tuning hyperparameters $\\eta, \\alpha$ seems like it could be difficult. The automatic $\\eta$ normalization scheme still likely requires hyperparameter tuning based on the variance of Q values.\n\n**Response:**  \n\nWe appreciate the reviewer’s insightful comment. It is true that RVDT introduces two additional regularization terms on top of the DT loss, which brings some extra hyperparameter tuning effort to ensure optimal performance. We summarize the loss in Equation (7) into the following form:\n$$L(RVDT) = L(DT) - \\eta L(Q) + \\alpha L(KL).$$\n\nIn practice, we find that certain tuning heuristics can effectively reduce the tuning burden, especially when considering the relative reliability of the two regularizers across different types of datasets. For instance, in environments where Q-value estimation is unreliable—such as those with sparse rewards—we typically set $\\eta$ to a relatively small value. This shifts the optimization emphasis toward the expert KL regularizer, which tends to provide a more stable and effective learning signal in such scenarios. Conversely, in environments with dense and reliable reward signals, a relatively larger $\\eta$ may be preferred to better leverage the Q-regularizer.\nAs a result, tuning $\\eta$ and $\\alpha$ is generally not prohibitively difficult.\n\nFrom a theoretical perspective, Proposition 1 in our paper provides partial justification for the compatibility of the two regularizers. Specifically, under the Gaussian policy assumption, the expert KL regularizer can be interpreted as implementing a weighted sampling scheme, leading to the following reformulated loss:\n$$ L_{RVDT}(θ) =  E_{τ \\sim D} [ (1 + α · I[τ ∈ D_e]) · ( Σ_{i=1}^H -log π_θ(a_i | s_i, g(τ_i), τ̄_{t-1}^K) ) ] - \\eta L(Q).$$\nThis formulation suggests that the expert KL regularizer effectively operates through a sampling-based mechanism that is inherently compatible with the main DT loss. As such, it does not introduce significant conflict during optimization and maintains stability in practice.\n\n---\n\n### **Question 1: Clarity of Contributions**\n\n> Contributions: Each contribution is a bit long and difficult to parse without reading the preliminary section. Could you simplify/summarize these a bit?\n\n**Response:**  \n\nWe appreciate the reviewer’s request. To improve clarity, we provide a simplified summary of our contributions below.\n\n1. We theoretically characterize how the performance of CSM-based policies depends on the coverage of high-quality (expert-level) and runtime (full-spectrum) returns in the offline dataset.\n2. We propose a simple yet effective return-coverage rebalancing mechanism, which can serve as a plug-in module to improve existing CSM-based algorithms.\n3. We design RVDT, a new CSM-based offline RL algorithm built upon DT, incorporating both Q-value and expert policy KL-divergence regularization for better alignment between the expected returns of sampled actions and the optimal returns.\n\nThe current presentation of our contributions in the paper prioritizes completeness, which is why we include relevant technical details to ensure precision and clarity in context.\n\n---\n\n### **Question 2: Intuition Behind the Conditioning Function in Section 3.2**\n\n> Section 3.2: Can you please add a sentence with a more intuitive explanation of the conditioning function $f$? The $f$ we aim to find is often the RTG of the optimal policy $\\pi^\\star$, correct? $f(s) = \\max_{\\pi}\\mathbb{E}_{\\pi}[g(s)]$\n\n**Response:**  \n\nThank you for the insightful comment, and your understanding is indeed correct. We will add a more intuitive explanation of the conditioning function $f$ to Section 4.1, immediately following the sentence: \"Our analysis extends this perspective by considering a more general function class that encompasses both near-optimal and potentially non-optimal runtime conditioning functions, denoted as $f^\\star$ and $f$, respectively.\" (Line 162). The added explanation will be as follows:\n\nThe runtime conditional return function $f$ corresponds to arbitrary possible returns that may be encountered during policy execution. Let $\\mathcal{G}$ denote the collection of all possible returns collected by $\\pi \\in \\Pi$, then $f: \\mathcal{S} \\rightarrow \\mathcal{G}$. The target conditional function $f^\\star$ that we aim to find is the RTG under the optimal policy $\\pi^\\star$, i.e., $f(s) = \\max_{\\pi}\\mathbb{E}_{\\pi}[g(s)]$.\n\n---\n\n### **Limitations: Limitations Not Discussed**\n\n> Limitations: I do not think the authors list any limitations. I would suggest they consider adding a few sentences about this. \n\n**Response:**  \n\nWe sincerely thank the reviewer for the helpful suggestion. We will add a discussion in our revision to state the limitations of the proposed RVDT algorithm. The content is as follows:\n\nThe proposed RVDT algorithm has several limitations:\n1. RVDT requires learning an additional imitation policy, which introduces extra computational overhead.\n2. RVDT is currently designed as a purely offline RL algorithm. How to effectively incorporate online interactions to expand return support remains an open question and will be a promising direction for future work.\n\n---\n\nWe sincerely appreciate the reviewer’s feedback and will incorporate these suggestions to improve clarity and completeness in the revised version, and we hope our response can alleviate some of your concerns."}]

Author rebuttal:

Meta-review: (a) Summary

The paper analyzes why conditional sequence modeling (CSM) policies in offline RL succeed or fail by linking performance to two kinds of return coverage—expert-level and full-spectrum—and derives bounds and propositions formalizing this dependency. Building on the analysis, it proposes RVDT, a DT-based objective combining expert-policy KL (explicit rebalancing) with Q-guided regularization (implicit rebalancing). Across D4RL (Gym, Kitchen, AntMaze, Maze2D), RVDT matches or exceeds strong CSM baselines and is notably better in low-data and sparse-reward settings.

(b) Strengths
- The paper clearly identifies and studies the impact of expert-level and full-spectrum return coverage on CSM performance, providing useful theory. [Lvw8][4BE1]
- The RVDT objective is simple to implement yet effective, yielding strong results and clean ablations. [Lvw8][Mnay]
- The experiments are broad and systematic, covering data quality, scale (low-data), and sparse-reward regimes. [Lvw8][4BE1]
- Writing and presentation are clear, making complex material accessible. [Mnay]

(c) Weaknesses
- Theoretical guarantees hinge on a return-coverage assumption that can be weak or vacuous when coverage is tiny. [8NTa]
- Gains over QT are modest on several Gym/Kitchen tasks, and the method is close in spirit to QT. [Mnay][8NTa]
- The method depends on selecting “expert” trajectories, which may be ambiguous in some real-world domains, and KL regularization does not expand support into unseen state-action regions. [4BE1]
- Extra compute is required (training an expert policy and two regularizers), raising cost-effectiveness questions. [Lvw8][8NTa]

(d) Discussion summary

Most of the questions and comments were satisfactorily addressed in the author rebuttal

(e) Decision rationale

The work offers (1) a clear theoretical account of return-coverage in CSM and a practical, plug-in rebalancing mechanism; (2) consistent improvements with especially strong margins in low-data and sparse-reward settings; and (3) thorough ablations and added analyses that addressed reviewers’ core questions.
Decision: Accept (poster)