Free from Bellman Completeness: Trajectory Stitching via Model-based Return-conditioned Supervised Learning
Authors: Zhaoyi Zhou
Chuning Zhu
Runlong Zhou
Qiwen Cui
Abhishek Gupta
Simon Shaolei Du
Venue: ICLR 2024 poster (2024)
Keywords: Offline Reinforcement Learning
Return-Conditioned Supervised Learning
Bellman Completeness
Trajectory Stitching
Primary area: reinforcement learning
TL;DR: We analyze advantage of return-conditioned supervised learning in near-deterministic environments and improves return-conditioned supervised learning to enable trajectory stitching.

Abstract: Off-policy dynamic programming (DP) techniques such as $Q$-learning have proven to be important in sequential decision-making problems. In the presence of function approximation, however, these techniques often diverge due to the absence of Bellman completeness in the function classes considered, a crucial condition for the success of DP-based methods. In this paper, we show how off-policy learning techniques based on return-conditioned supervised learning (RCSL) are able to circumvent these challenges of Bellman completeness, converging under significantly more relaxed assumptions inherited from supervised learning. We prove there exists a natural environment in which if one uses two-layer multilayer perceptron as the function approximator, the layer width needs to grow *linearly* with the state space size to satisfy Bellman completeness while a constant layer width is enough for RCSL. These findings take a step towards explaining the superior empirical performance of RCSL methods compared to DP-based methods in environments with near-optimal datasets. Furthermore, in order to learn from sub-optimal datasets, we propose a simple framework called MBRCSL, granting RCSL methods the ability of dynamic programming to stitch together segments from distinct trajectories. MBRCSL leverages learned dynamics models and forward sampling to accomplish trajectory stitching while avoiding the need for Bellman completeness that plagues all dynamic programming algorithms. We propose both theoretical analysis and experimental evaluation to back these claims, outperforming state-of-the-art model-free and model-based offline RL algorithms across several simulated robotics problems.

Published as a conference paper at ICLR 2024
FREE FROM BELLMAN COMPLETENESS: TRAJECTORY STITCHING VIA MODEL-BASED RETURN-CONDITIONEN SUPERVISED LEARNING
Zhaoyi Zhou', Chuning Zhu?, Runlong Zhou', Qiwen Cui', Abhishek Gupta', Simon S. Du²
Institute for Interdisciplinary Information Sciences, Tsinghua University
2 Paul G. Allen School of Computer Science and Engineering, University of Washington
zhouzhac220maills weaiiiss.ssssgge tsinguua. edu. en Tzchuning, vectorzh, qwcui, abhgupta, ssdu (cs. washington.
ABSTRACT
Off-policy dynamic programming (DP) techniques such as Q-learning have proven to be important in sequential decision-making problems. In the presence of function approximation, however, these techniques often diverge due to the absence of Bellman completeness in the function classes considered, for the success of DP-based methods. this paper, we show how off-policy learning techniques based on relurn-conditioned supervised learning (RCSL) are able to circumvent these challenges of Bellman completeness, converging under significantly more relaxed assumptions inherited from supervised learning. We prove there exists a natural environment which if one uses two-layer multilayer perceptron the function approximator, the layer width needs to grow linearly with the state space size to satisfy Bellman completeness while a constant layer width is enough for RCSL. These findings take step towards explaining the superior empirical performance of RCSL methods compared to DP-based methods in environments with near-optimal datasets. Furthermore, in order to learn from sub-optimal datasets, we propose simple framework called MBRCSL, granting RCSL methods the ability of dynamic programming to stitch together segments from distinct trajectories. MBRCSL. leverages learned dynamics models and forward sampling to accomplish trajectory stitching while avoiding the need for Bellman completeness that plagues all dynamic programming algorithms. We propose both theoretical analysis and experimental evaluation to back these claims, outperforming -fate-of-the-art model-free and model-based offline RL algorithms across several simulated robotics problems
crucial condition
INTRODUCTION
The literature reinforcement learning (RL) has yielded a promising set of techniques for sequential decision-making making, with several results showing the ability to synthesize complex behaviors in variety of domains (Mnih et al. 2013 Lillicrap et al., 2015 Haarnoja et al.) 2018) even in the absence of known models of dynamics and reward in the environment. Of particular interest in this work is class of reinforcement learning algorithms known as off-policy RL algorithms (Fujimoto et al.) 2018). Off-policy RL algorithms are those that are able to learn optimal behavior from "off-policy" dataset i.e. a potentially suboptimal dataset from some other policy. Moreover, off-policy RL methods typically able to "stitch" together trajeeerctsoies across many different collecting trajectories, thereby showing the ability to synthesize behavior that better than any trajectories present in the dataset (Degris al. 2012 Kumar al. 2019a).
The predominant class of off-policy RL algorithms builds on the idea of dynamic programming (DP) (Bellman) 1954 1957) learn state-action Q-functions. These techniques iteratively update the Q-function by applying the Bellman operator on the current Q-function. The step-wise update of
'Our code is available at https: httptttttttttt/tt/////////ttottt ccm iileoyiioouu 123 233mbrrs) Part of the work was
done while Zhanyi Zhou was visiting the University of Washington.

---

Published as a conference paper at ICLR 2024
Behavior Policy
RCSL
Offline Data
Model Rollouts
Rollout Data
Deturn-Conditione Policy
Dynamics Model
Figure I: Model-based renurn-conditioned supervised learning (MBRCSL). Model-based rollouts augments datset with potentially optimal trajectories, enabling trajectory stitching for RCSL while avoiding Bellman completeness requirements.
DP enables the advantages listed above usage of prior suboptimal data and synthesis of optimal policies by stitching together multiple suboptimal trajectories. On the other hand, these off off-policy RL algorithmss based on dynamic programming can only guarantee convergence under the condition of Bellman completeness, a strong and nuanced claim on function approximation that ensures the projection error during iterative DP can be reduced to be arbitrarily low. Unlike supervised learning, designing function class that satisfies Bellman completeness is particularly challenging, simply enlarging the function class does not guarantee Bellman completeness.
This begs the question the requirements of Bellman completeness, while retaining the advantages of DP-based off-policy RL? To this end, study class of off- off-policy RL algorithms referred to cetur-conditioned supervised learning (RCSL) (Schmidhuber) 2019) (Srivastava et al.) 2019). The key idea of RCSL to learn return-conditioned distribution actions in each state by directly applying supervised learning on the trajectories in the dataset, achieving satisfactory performance by simply conditioning the policy high desired returns. These techniques have seen recent prominence due to their surprisingly good empirical performance. simplicity and stability (Kumar et al. 2019b). While RCSL has empirically shown good results, the reason behind these successes in many empirical domains is still not well understood. In this work, we show how RCSL can circumvent the requirement for Bellman completeness. We theoretically analyze how RCSL benefits from not requiring Bellman completeness and thus can provably outperform DP-based algorithms with sufficient data coverage.
design off off-policy reinforcement learning algorithms that are free from
While RCSL able to circumvent the requirement for Bellman completeness, comes with the requirement for data coverage. As prior work has shown, RCSL has fundamental inability to stitch trajectories (Brandfonbrener et al. 2022), a natural advantage of using DP-based off-policy RL. As long the dataset does not cover optimal trajectories, RCSL methods are not guaranteed to output the optimal policy. To remedy this shortcoming, while preserving the benefit of avoiding Bellman completeness. we propose novel algorithm- modell-asee RCSL (MBRCSL, cf. Figure T.. The key insight in MBRCSL is that RCSL style methods simply need data coverage, and can take care of optimizing for optimal behavior when the data provided has good (albeit skewed) coverage. Under the Markovian assumption that is made in RL, we show data coverage can be achieved by simply sampling from the behavior policy under learned model of dynamics. We show how conjunction of two supervised learning problems learning model of dynamics and of the behavior policy can provide us the data coverage needed for the optimality of RCSL. This leads to method that both avoids Bellman completeness and is able to perform stitching and obtain the benefits of DP.
Concretely, the contributions of this work are:
1. We theoretically show how RCSL can outperform DP when Bellman completeness is not satisfied. We theoretically show the shortcomings of RCSL when the off-policy dataset does not cover optimal trajectories.
3. We propose a novel technique MBRCSL that can achieve trajectory stitching without suffering from the challenge of Bellman completeness.
We empirically show how MBRCSL can outperform both RCSL and ba-ba off-policy RL algorithms in several domains in simulation.
2

---

Published as a conference paper at ICLR 2024
1.1 RELATED WORK
Function Approximation and Bellman Completeness. Function approximation is required when the state-action space is large. A line of works has established that even if the optimal Q-function is
realizable, the agent still needs exponential number of samples (Du et al., 2019 Wang et al.) 2020
202T; Weisz et al. 2021). Bellman completeness is a stronger condition than realizability but
permits sample-eficiinn algorithms and has been widely adopted in theoretical studies (Munos &
Szepesváni 2008 Chen & Jiang, 2019 Jin et al. 2020 Zanette et al., 2020 Xie et al.) 2021).
Return-conditioned supervised learning. Return- centin-condi supervised learning (RCSL) directly
trains a return-conditioned policy via maximum likelihood (Schmidhuberi 2019 Srivastava et al.
2019; Kumar et al. 2019b; Andrychowicz et al. 2017 Furuta et al., 2021; Ghosh et al. 2019))
The policy architecture ranges from MLP (Emmons et al., 2022] Brandfonbrener et al. (2022) to
transformer (Chen et al. 202I Bhargava et al.) 2023) diffusion models (Ajay et al. 2022).
The finding that RCSL cannot do stitching is first proposed by (Kumar et al.)
2022) through empirical studies. Here, the trajectory stitching refers combination of transitions from different trajectories to form potentially better trajectories, which is a key challenge to offline RL algorithms (Char et al.,
2022 Hepburn& & Montana 2022; Wu et al. 2023). BBranfoonrener et al. (2022) gives an intuitive counter-example, but without a rigorous theorem. In Section 3.2, we provide two different reasons and rigorously show the the inability of RCSL to do stiching. BBrandfonrrener et al. (2022) also derives sample complexity bound of RCSL. However, they do not concretely characterize when RCSL requires much weaker conditions than Q-learning find the optimal policy. |Wa al. (2023): ;
Liu & Abbeel )2023): Yamagata et al. (2023) try to improve RCSL performance on sub-optimal trajectories by replacing dataset returns with (higher) estimated returns, while our approach aims
at extracting optimal policy from dataset through explicit trajectory stitching. Besides, Paster et al. (2022): Yang et al. (2022) address the problems of RCSL in stochastic environments.
Model-based RL. Model-based RL learns approximate dynamics model of the environment and extracts a policy using model rollouts (Chua et al., 2018: Janner et al. 2019: Hafner et al. 2021 Rybkin et al.] 2021). A complementary line of work applies model-based RL to offline datasets by incorporating pessimism into model learning or policy extraction (Kidambi et al. 2020 Yu et al.. 2020 2021). ). We leverage the fact that models do not suffer from Bellman completeness and use it as a crucial tool to achieve trajectory stitching.
2 PRELIMINARIES
Finite-horizon MDPs. A finite-horizon Markov decision process (MDP) can be described as M (H, S, A, μs, T, r). where H is the planning horizon, S is the state space, and A is the action space. /4 E (SS) is the initial state distribution, and T : S x A + A(S) is the transition dynamics. Note that for any set X, we use A(X) to denote the probability simplex over X. In case of deterministic environments, we abuse the notation and write T S A + S. S x A + [0.
Trajectories. A trajectory is a tuple of states, actions and rewards from step 1 to H: T ($1, $1,44.77
SH, az rH). Define the return-to-go (RTG) of trajectory at timestep h as 9h = Σ4 Σm'= rh. The alternative representation of trajectory is r = (81, 91, 01; 38H: 9H, aH). We abuse the notation and call the sub-trajectory T/i : j] = (si az: ;8 Gj. az) as trajectory. Finally, we denote T as the set of all possible trajectories.
Policies. A policy as input. If the trajectory contains RTG yh, then the policy is a relurn-conditoned policy. The Markovian property of MDPs ensures that the optimal policy falls in the class of Markovian policies, IIM = { : - (. }. A policy class with our special interest is Markovian return- cruur-conditioned policies. MRM { S I + A(A)), i.e., the class of one-step conditiond policies.
T + A(A) is the distribution of actions conditioned on the trajectory
Value and Q-functions. We define value and Q-functions for policies in 11M, Let IE, [-] denote the expectation with respect to the random trajectory induced by MD in the MDP M, that is, 7 ($1, a1, Ti...., SH aH, TH), where $1 ~ pa(-), an ~ (-(ss)). yy = r (sh. an), 8h+1 ~ T T(- nn) for any h E [H]. We define the value and Q-functions using the RTG gh:
VA (s) Ex [gt shh
s], Q7 (s. a) Ez [g s, an a)
3

---

Published as a conference paper at ICLR 2024
Algorithm 1 Return-conditioned evaluation process
1: Input: return-conditioned policy S x R + (AA), desired RTG ÿ1-
2: Observe 81-
3: for h = 1,2, H do
4: Sample action on ~ (-118h> 14h-0i Ih)-
5: Observe reward Th and next state sh+1-
6: Update RTG 9h+11 9h Th.
7: Return: Actual RTG g Em=1 Th-
For convenience, we also define + O, 1 = 0, for any policy т and 8, a.
Expected RTG. Return-conditioned policies needs an initial desired RTG ÿI to execute. The return- conditioned evaluation process for return-come conuin-onddiooeed policies is formulated Algorithm I We know that 91 is not always equal g, the real RTG. Hence to evaluate the performance of MRC. we denote by [TT 91) the expected return under æ with initial desired RTG 91-
.*). where r' = [s] aj ri. ... Sy fr rH) (i E [K]) is a trajectory sampled from the environment. We abuse the notation and define (s,
Offline RL. In offline RL, we only have access some fixed dataset D (T),
D if and only if there exists some i [K] and h [H] such that (s%, 4;) = (s,a). a). We say dataset DD uniformly covers M if (s, a) E D for any E S, a E A. For A of M, say dataset 'D covers nt if (s, (s)))
any deterministic optimal policy x² S
'D
for any E S. The goal of offline RL is to derive a policy using 'D which maximizes the expected total reward. For Markovian policy π E I1M, the goal is g max, E. (T)) For Markovian return-conditioned policy क E [MMRC the goal is max, (
2.1 REALIZABILITY AND BELLMAN COMPLETENESS OF O-LEARNING
Many exisiting RL algorithms, such Q-learning and actor-critic, apply dynamii-programming (DP) technique (Bertsekas & Tsitsiklis 1996) to estimate value functions or Q-functions. For instance,
Q-learning updates estimated (optimal) Q-function Q (8, a) by
Qn BQn-1 (r(s,a) Bx-T((sa sup Qu+1 (s', a') d')
(I)
a'
where B is called the Bellman operator (Sutton & Barto 2018).
In practice, @ is represented using some function approximation class such multilayer perceptron (MLP). Two crucial conditions for DP-based methods with function approximation to succeed are realizability (Definition D and to be Bellman complete (Definition) 2 (Munos 2005 Zanette 2023).
Definition (Realizability). The function class F contains the optimal Q-function: Q" F.
Definition 2 (Bellman complete). A function approximation class F is Bellman complete under Bellman operator B if maxqer minq'eF ||Q' - O.
Remark 1. It has been shown that realizability alone is not enough, as Weisz et al. (2021) Wang et al. 2020 show that it requires an exponential number of samples to find a near-optimal policy.
2.2 NETURN-CONDITIONE SUPERVISED LEARNING
RCSL methods learn a retum-condittoned policy -cx via supervised learning, where cx is the policy context. We focus on return-conditioned policies with context 1, i.e., policy of the form n! (ax [sh, gh), because the optimal policy can be represented by eetur-conditioned policy of context 1. Next, we introduce the RTG dataset, summarizes the information that RCSL algorithm relies on.
1], Sh. Ih from D
Definition 3 (RTG dataset). Given a collection of trajectories D ( , an RTG dataset for D is defined by D := =( , | sts
We are now ready to define the RCSL framework.
Definition (RCSL framework). An offline RL algorithm belongs to RCSL.framework if it takes an RTG dataset input and outputs a return-conditoondd policy.
4

---

Published as a conference paper at ICLR 2024
3 UNDERSTANDING THE STRENGTHS AND WEAKNESSES OF DETTUR-CONNITIONSS SUPERVISED LEARNING METHODS
In this section, first provide theoretical discussion of why RCSL methods can outperform the DP methods in deterministic environments, when there is sufficient data coverage. We then give rigorous theoretical treatments explain why RCSL style methods cannot do "rajjcttry--stichinng.
3.1 INSIGHT 1: RCSL CAN OUTPRRFORM DP IN DETERMINISTIC ENVIRONMENTS
In this section, show that RCSL can achieve better performance than DP-based algorithms, as long as the dataset contains expert (optimal) trajectory. This stems from the fact that DP-algorithms needs satisfy Bellman-complleennss when function approximation is used, while RCSL algorithms only need to represent the optimal policy, and hence only require realizability.
To explain this point, choose representative implementation from both RCSL and DP-bases methods respectively to make a comparison. For the former, we design "MLP-RCSL", a simple RCSL algorithm that learns deterministic policy. Suppose that the action space A is a subset of Rd. MI.P-RCSL trains policy network (8 g) E IR by minimizing mean square error (MSE) of action (or more generally, maximize the log-likelihod):
[(0) = ()... || (8,9) - all².
(2)
We add projection f(a) arg min ||a ( at the output of policy network, so that we are ensured get a valid action.
We choose Q-learning as a representative for DP-based methods, where the Q-function is learned by Qnnetwork, denoted by Qs(s, a). To make the comparison fair, both 7s(s, and Q0(s, a) are implemented with a two-layer MLP network with ReLU activation.
Often DP-based offline RL methods aim to penalize out-of out-of- distribution actions (in the presence of incomplete coverage) via techniques such as conservativism (Kumar et al. (020) or constrained optimization (Wu el al. 2019). However, we do not need to consider these techniques in our constructed environment, conservatism is not needed when the dataset has full coverage.
Remark 2.
We first analyze when MLP-RCSL and Q-learning satisfy their own necessary requirements to succeed, terms of the hidden layer size needed. To be specific, we set the requirements as:
MLP-RCSL: The policy network must be able to represent the optimal policy.
Q-learning: The Q-network must be able to represent the optimal Q-function and satisfy Bellman completeness.
Theorem 1. There exists series of MDPs and associated datasets {M" D (i) S O(u);
(H",S") A.A u, "", Juz1. such that the following properties hold:
(ii) D uniformly covers M";
(iii) There exists a series of MLP-RCSL policy network (two-layer MLP) { Juz1 such that
TRC represent optimal policy for M with O(1) neurons in the hidden layer;
(iv) For u 2 1, any two-layer MLP Q-network representing the optimal Q-function and satisfying Bellman completeness simultaneously should have (2()) neurons in the hidden layer.
We also provide simulations to verify the theorem (cf. (cI.
Remark 3. From Theorem 7 we that the model size of Q-leurning must grow linearly with state space. On the other hand, RCSL can scale with large stute space Besides, we also point out that the success of RCSL is based the existence of expert trajectory.
Proof idea. Detailed proof of Theorem I is deferred to Appendix A We construct a class of MDPs called LinearQ (Figure 2 such that the optimal policy as well as optimal Q-function can be represented by the linear combination of constant number of ReLU neurons, while the reward function must be represented with 12((S*)) ReLU neurons because there are that many. We denote the all-zero MIP as Q which satisfies Qotor (s, a) = 0 for all E SH and a E A. We see that Bellman operation (cf. Equation I applies to the all-zero MLP give (0) a) = (s, Vs S" A. Thus Q((8"l) hidden neurons needed to satisfy Bellman completeness.
5

---

Published as a conference paper at ICLR 2024
1
2
M
u
O-tunction
Q = 0)
Qls,
3u +
20-
2u
2n
i +
ris, 0)
3u + 2
reward
10
r(s, a 1]
2u
2u
3m
3u
state
Figure 2: Left: Illustration for LinearQ transition dynamics. This figure is valid for u being an even number. The transitions are plotted in red and blue arrows, where red arrows stand for 0 and blue ones for 1. The rewards are omitted. Upper right: Illustration for optimal Q-function of LinearQ with respect states, where size parameter neurons. Lower right: Illustration for reward function of LinearQ with respect to states, where size parameter u 12((5s)) hidden neurons are needed represent reward.
4. Realizability can be satisfied with constant hidden
3.2
INSIGHT 2: RCSL CANNOT PERFORM TRAJECTORY-STIICHIIE
is generally believed that RCSL, cannot perform trajectory stitching (Brandfonbrener al. 2022 Wu et al. 2023). We provide two theorems to rigorously show the sub-optimality of RCSL algorithms, even when the environment has deterministic transition and dataset satisfies uniform coverage. These two theorems present different rationales for the failure of trajectory stitching. Theorem shows that Markovian RCSL (context length equal to 1) outputs a policy with a constant sub-optimal gap. Theorem 3 states that any decision transformer (using cross-entropy loss in RCSL) with any context length fails to stitch trajectories with non-zero probability.
Theorem For any Markovian RCSL algorithm, there always exists an MDP M and an (sub- optimal) offline dutaset 'D such that
(i) M has deterministic transitions;
(ii) D uniformly covers M;
(iii) The output return-conditioned policy satisfies - J=x, 2 1/2 almost surely.
The key idea is that the information of reward function cannot be recovered from the RTG dataset D We are able to construct two MDPs with different reward functions but the same RTG dataset. Therefore, at least RCSL will least fail on one of them. The proof is deferred Appendix B.]
Theorem 3. Under the assumption that decision transformer outputs a mixture of policies for trajectories in the dataset for an unseen trajectory (details in Assumption 7), there exists an MDP M and an (sub-optimal) offline dataset such that
(i) M has deterministic transitions;
(ii) D uniformly covers M;
(iii) Any decision transformer æ satisfies J(T,g') <
The key idea is that cross-entropy loss of decision transformers forces non-zero probability for imitating trajectories in the dataset. We construct dataset containing two sub-optimal trajectories with the same total reward, then their first step information ($1.91) are the same. We make their first step action an and a' different, the RTG dataset contains trajectories (81, 91, a1) and ($1, 91, af). Under the assumption on generalization ability to unseen trajectory (Assumption 1.. the decision transformer will randomly pick actions with the first step (s1,g" > g1).
MODEL-BASED RETURN- CONDRI-COND SUPERVISED LEARNING
In this section, we ask whether we can avoid dellman-completter requirements, while still being able to stitch together sub-optimal data. To this end, propose model-based returncoodiitooedd

---

Published as a conference paper at ICLR 2024
Algorithm 2 MBRCSL: Model-Based Return-Conditioned Supervised Learning
Require:
Offline dataset D, dynamics model To. behavior policy My and output return- policy Tg. desired rollout dataset size n.
1: Collect initial states in D to form a dataset Dinitial. Compute the highest return Jmax in D.
2: Train dynamics model To a) and behavior policy (a(a)s) on D via MLE.
3: Initialize rollout dataset Pollont f Ø, rollout dataset size counter cnt < 0.
4: for = 1,2, do
5: Rollout a complete trajectory T from ps and Te, with initial state sampled from Dinitial-
6: Compute the predicted total return g.
7: if g >> Imax then
Add T to D rollout- cnt <- cnl 1. 10: if cnl = then break
8:
9:
11: Train To (a s, g) on Proloout via MLE.
supervised learning (MBRCSL) (ef. Algorithm] 2). This technique brings the benefits of dynamic programming to RCSL without ever having to do dynamic programming backups, just forward simulation. As demonstrated in Section 3.1 when there is sufficient data coverage, RCSL can outperform DP, but typically this data coverage assumption does not hold.
The key idea behind MBRCSL is to augment the offline dataset with trajectories that themselves combine sub-components of many different trajectories (i.e "stitching") "), on which we can then apply RCSL to acquire performance that is better than the original dataset. Since we are in the off-policy setting and want to avoid the requirement for additional environment access, we can instead learn an approximate dynamics model from the off-policy dataset, well an expressive multimodal behavior policy (Line 2 both using standard maximum likelihood estimation (MLE). Then, by simply rolling out (sampling i.i.d and executing sequentially) the behavior policy on the approximate dynamics model, we can generate "stitched" trajectories (Lines 4 10). In many cases, this generated data provides coverage over potentially optimal trajectory and return data, thereby enabling RCSL to accomplish better- than-tth---aaaa dataset performance. We aggregate optimal trajectories into new rollout dataset by picking out trajectories with returns larger than the maximum return in the dataset. Finally, applying RCSL on the augmented rollout dataset enables trajectory stitching without the need to satisfy Bellman completeness (Line TTp. Another advantage of MBRCSL framework is its modular architecture, i.e., dynamics model, behavior policy and output policy architecture can be designed independently. This allows us to leverage the most performant model classes for each individual sub-component. The precise details of the architectures and optimization algorithms that empirically work best in each domain are left to the experiments (Section
5 EXPERIMENTS
In our experimental evaluation, we are primarily focused on evaluating (a) how MBRCSL performs on offline RL tasks with sub-optimal datasets, comparison with model-free and model-based offline RL baselines? Moreover, we hope to understand (b) How each component in MBRCSL affects overall performance? Question (b) is investigated in Appendix with a complete ablation study.
5.1 EVALUATION ON POINT MAZE ENVIRONMENTS
We first evaluate the methods custom Point Maze environment built on top of D4RL Maze (Fu et al. 2020). As illustrated in Fig. 3 this task requires the agent to navigate ball from the initial location (denoted S) designated goal (denoted G). To investigate if the methods can do stitching.
we construct an offline dataset consisting of two kinds of suboptimal trajectories with equal number:
A detour trajectory + A + B + G that reaches the goal in a suboptimal manner.
A trajectory for stitching: S M. This trajectory has very low return, but is essential for getting
the optimal policy.
The optimal trajectory should be stitching of the two trajectories in dataset (S + M G). The trajectories are generated by a scripted policy similar to that in D4RL dataset (Fu et al. 2020 The resulting dataset has averaged return 40.7 and hightest return 71.8.
7

---

Published as a conference paper at ICLR 2024
Figure 3: Point Maze illustration. Left and middle: Simulation images initial and final timestep. Green point represents the ball to control, red point represents the target position. Right: Maze and dataset description. The initial point is S and the goal G. The dataset consists of two trajectories: S + A B G and S + M.
To answer question (a), compare MBRCSL with several offline RL baselines. (1) CQL (Kumar ct al. 2020) a model-free offline RL method that learns a conservative Q function penalizing high values for out of distribution actions. (2) MOPO Yu al. (020) is a model-based offline RL method that learns an ensemble dynamics model and performs actor-critic learning using short-horizon model-rollouts. (3) COMBO (Yu et al. 2021) combines CQL and MOPO by jointly learning conservative Q function and a dynamics model. (4) MOReL (Kidambi et al., 2020) first learns a pessimistic MDP from dataset and then learns near-optimal policy in the pessimistic MDP. (5) Decision Transformer (DT) (Chen et al. 2021) is RCSL method that leverages transformer model generate return-conding condittood trajectories. (6) %BC performs behavior cloning on trajectories from the dataset with top 10% return.
We implement MBRCSL with the following architectures:
Dynamics model learns a Gaussian distribution over the next state and reward: (o re|s;, az)
N' ((8++1, (8,, (8),)), Eg( model is trained independently via maximum- likelliooi.eiiod.. This is the same as many existing model-based approaches such MOPO (Yu al. (020)) and COMBO (Yu al. 2021))
a.)). We learn ensemble of dynamics models, in which each
We learn an approximate behavior policy with a diffusion policy architecture (Chi et al. 2023),
which is able capture multimodal action distributions. We fit the noise model by minimizing
mean squared eтror (MSE).
The return- ceturn-condiiions policy represented by a MLP network and trained by minimizing MSE.
As shown in [ MBRCSL outperforms all baselines on the Point Maze environment, successfully stitching together suboptimal trajectories. COMBO, MOPO, MOReL and CQL. struggles to recover the optimal trajectory, potentially due to challenges with Bellman completeness. They also have large variance in performance. addition, we show in Appendix D that performance of DP-based offline RL methods cannot be improved by simply increasing model size. DT and %BC have low variance, but they cannot achieve higher returns than that of datset due to failure to perform stitching.
Table Results on Point Maze. We record the mean and standard deviation (STD) on 4 seeds. The dataset has averaged return 40.7 and highest return 71.8.
Task
MBRCSL (ours)
COMBO
MOPO
MOReL
CQL
DT
BBC
Pointmaze
91.5
56.6450.1
77.9
26.4428.1
34.8 24.9
57.244.
54,0+9.2
5.2 EVALUATION ON SIMULATED ROBOTICS ENVIRONMENTS
We also evaluate the methods on three simulated robotic tasks from Singh al. (2020) In each task, a robot arm is required to complete some goal by accomplishing two phases of actions, specified as follows:
(PickPlace) (1) Pick an object from table; (2) Place the object into tray.
CloseeDrrawer) (1) Open a drawer; (2) Grasp the object the drawer.
(BlockedDrawer) (1) Close drawer top of the target drawer, and then open the target drawer; (2) Grasp the object in the target drawer.
We take images from Singh al. (2020) for environment illustration (cf. Figure 4 the associated dataset consists of two kinds of trajectories: (i) Prior trajectories that only perform phase 1; (ii) Task trajectories which only perform phase 2, starting from the condition that phase is completed. For PickPlace, ClosedDrawer and BlockedDrawer, the prior dataset completes phase with probability about 40%, 70% and 90% respectively. The task dataset has success rate of about
For each task,
8

---

Published as a conference paper at ICLR 2024
90% for PickPlace and 60% for both ClosedDrawer and BlockedDrawer. The agent must stitch the prior and task trajectory together to complete the task. All three tasks have sparse reward, with return i for completing phase 2, and 0 otherwise. Comparing to Point Maze, the simulated robotic tasks require more precise dynamics estimates.
We choose CQL (Kumar et al. 2020), COMBO (Yu et al., 2021) 1) and Decision Transformer (DT) (Chen al. 2021) as baselines of model- free, model-based and RCSL methods, respec- tively. To demonstrate performance improve- ment of the output policy over the behavior policy in MBRCSI, we also tested behavior cloning implemented with diffusion policy (de- noted by BC), which has the same architecture as the behavior policy in MBRCSL.
To capture the more complicated dynamics, we implemented the dynamics and output return- conditioned policy with the following architec- tures:
Figure 4: Simulated robotic tasks illustrations from CoG (Singh al. 2020). The rows represent Pick- Place, ClosedDrawer and BlockedDrawer, respec- tively, from top to bottom. The dataset for each task consists of two kinds of trajectories: Some perform actions in red arrows and others perform actions in green arrows. A task is completed if and only if actions in both red and green arrows are performed successfully.
Dynamics model is represented with transformer model (Vaswani et al.) (017)) Te(s', rrss, a). which trained to minimize cross-entropy loss.
The output return-conditioned policy is rep- resented by an autoregressive model (Ko- renkevych al. 2019) trained by MILE.
As shown in Table 2 MBRCSL outperforms all baseline methods
three tasks. CQL
achieves nonzero while low success rates and suffers from high variance in two of the tasks. This is possibly because Bellman completeness and sparse reward impede correct Q-function estimation. COMBO fails to complete the task all, potentially due to inaccurate model rollouts incurring a higher Q estimation error. DT fails to extract meaningful behavior because of lack of expert trajectory. We also found that compared to the behavior policy (BC) which tries capture all possible behaviors, the output policy (MBRCSL) is able to extract successful trajectories from rollouts and results in higher success rate.
Table 2: Results on simulated robotic tasks. We record the mean and STD of the success rate seeds.
Task
MBRCSL (ours)
CQL
COMBO
DT
BC
PickPlace
0.48
0.22+50.35
0+00
0+:0
0.07 0.03
ClosedDrawer
0.51
0.11+0.08
0::0
0+:0
0.38+:0.02
BrowkedDrawer
0.68
0.34
01:0
0.:0
0.61
CONCLUSION
In this work, we conduct a thorough analysis of off-policy reinforcement learning, theoretically showing how return-conditioned supervised learning algorithms can provide benefit over typical dynamic programming-bast based methods for reinforcement learning under function approximation through the viewpoint of Bellman completeness. We show that with data coverage, RCSL style algorithms do not require Bellman completeness, simply the easier realizability condition. We then characterize the limitations of RCSL in its inability to accomplish petter-than-aatee through trajectory stitching. To remedy this, while still retaining the benefit of avoiding Bellman completeness, we propose a novel algorithm- MBRCSL, that is able to accomplish data coverage through i.i.d forward sampling using a learned dynamics model and show its benefits. Notable open problems here concern how to make MBRCSL work in stochastic environments and scale these algorithms up more complex and large-scale environments.
behavior

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
SSD acknowledges the support of NSF IIS 2110170, NSF DMS 2134106, NSP CCF 2212261, NSF IIS 2143493, NSF CCF 2019844, NSF IIS 2229881. CZ is supported by the UW-Amazon Fellowship.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision- decision-making? Preprint arXiv:2211.15657, 2022.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. Advances in neural information processing systems, 30, 2017.
Richard Bellman. The theory of dynamic programming. Bulletin of the American Mathematical Society, 6036)55 1954.
Richard Bellman. markovian decision process. Journal of mathematics and mechanics, pp- 679-684, 1957.
Dimitri Bertsekas and John N Tsitsiklis. Neuro-dynamic programming. Athena Scientific, 1996.
Prajjwal Bhargava, Rohan Chitnis, Alborz Geramifard, Shagun Sodhani, and Amy Zhang, Sequence modeling is a robust contender for offline reinforcement learning. Preprint arXiv:2305. 2023.
David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does crtur-connitioneed supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 15:1542-1553, 2022.
Ian Char, Viraj Mehta, Adam Villaflor, John M Dolan, and Jeff' Schneider. Bats: Best action trajectory stitching. preprint arXiv:2204. 20.p222244.222., /2026, 2022.
Jinglin Chen and Nan Jiang. information-thhoreti considerations in batch reinforcement learning. In International Conference Machine Learning, pp. 1042-1051. PMLR, 2019.
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, 34:1 15084-15097 2021.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. Preprint arXiv:2303.0411 04.3030..137, 2023.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, pp. 4765, 2018.
Thomas Degris, Martha White, and Richard S Sutton. Off-policy actor-critic. arXiv: 4839, 2012.
Preprint
Simon Du, Sham Kakade, Ruosong Wang, and Lin Yang. Is good representation sufficient for sample efficient reinforcement learning? Preprint arXhv:1910.03016. 2019.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning?, 2022.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. Preprint arXiv: 0,0900, 2018.
10

---

Published as a conference paper at ICLR 2024
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. Preprint arXiv:2111.10364 2021.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. Preprint arXiv:1912.06088, 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. Preprint arXiv: arXXiv/801.000 1801.01200, 2018.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL httt:tttts///ppeeeeeett
Charles A Hephurn and Giovanni Montana. Model-based trajectory stitching for improved offline reinforcement learning. preprint arXiv:2211.11603 2022.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelz- imer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neu- ral Information Processing Systems 32: Anmual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 12598-1259,, 2019. URL htttps:://proceeednngg neetrseeeddnnnse.ee.p.p...... html
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. Conference on Leurning Theory, pp. 2137-2143. PMLR, 2020.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Preprint arXiv:2005.05951, 2020.
Dmytro Korenkevych, A Rupam Mahmood, Gautham Vasan, and James Bergstra, Autoregressive policies for continuous control deep reinforcement learning. Preprint 1rX1v:1903.11524, 2019.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems, pp. 11761-11771, 2019a.
Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-condiiioned policies. arXiv:1912.13465. 2019b.
Preprint
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 33:1 179-11991.1 2020.
Aviral Kumar, Joey Hong, Anikait Singh, and Sergey Levine. When should prefer offline reinforcement learning over behavioral cloning? preprint arXiv:2204.05611. 2022.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. Preprint arXiv: (St.::550909971 2015.
Hao Liu and Picter Abbeel. Emergent agentic transformer from chain of hindsight experience. Preprint arXiv:2305. 16554, 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013.
Rémi Munos. Error bounds for approximate value iteration. In AAAI Conference on Artificial Intelli- gence, 2005. URL https:::s:: //api. semant damanticacen icscholarcr wtt.e//ho.....roooooooooooaaaaaa

---

Published as a conference paper at ICLR 2024
Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research, 9(5), 2008.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count luck: Why decision transformers and rvs fail in stochastic environments. Advances in Neural Information Processing Systems, 35: 2022.
Oleh Rybkin, Chuning Zhu, Anusha Nagabandi, Kostas Danilidiss Igor Mordatch, and Sergey Levine. Model- reinforcement learning via latent-space collocation. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning, ICMI. 202 I, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 9190-9201. PMLR, 2021. URL http: htttp://proceedlnggg mlr. wptpss//w3//yybiin///..
html
Juergen Schmidhuber, Reinforcement learning upside down: Don't t predict rewards-just map them to actions. Preprint 1912: 2019.
Avi Singh, Albert Yu, Jonathan Yang, Jesse Zhang, Aviral Kumar, and Sergey Levine. Cog: Connect- ing new skills past experience with offline reinforcement learning. Preprint arXiv:2010.14500, 2020.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaskowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. Preprint arXiv: 1912.02877, 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Ruosong Wang, Dean P Foster, and Sham M Kakade. What are the statistical limits of offline rl with linear function approximation? Preprûnt arXiv:2010. 18995 2020.
Yuanhao Wang, Ruosong Wang, and Sham Kakade. exponential lower bound for linearly realizable mdp with constant suboptimality gap. Advances in Neural Information Processing Systems, 34:9521-9553 9533, 2021.
Gellért Weisz, Philip Amortila, and Csaba Szepesvári. Exponential lower bounds for planning in mdps with Ehearly-realizable optimal action-value functions. In Algorithmic Learning Theory, pp. 1237-1264. PMLR, 2021.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. Preprint arXiv:1911.11361 2019.
Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. preprint arXiv:2307.02484, 2023.
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistemn pessimism for offline reinforcement learning. Advances in neural information processing systems, 34:6683-6694, 2021.
Taku Yamagata, Ahmed Khalil, and Raul Montoo-Rodrriuez. Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl. International Conference Machine Learning, pp. 18989-39000 PMLR, 2023.
Sherry Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Separating what you can control from what you cannot. In The Eleventh International Conference Learning Representations, 2022.
Tianhe Yu. Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Preprint arXiv.2005.13239 2020.
12

---

Published as a conference paper at ICLR 2024
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 28954-28900
Curran Associates, Inc., 2021.
Andrea Zanette. When is realizability sufficient for off-policy reinforcement learning?, 2023.
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill, Learning near optimal policies with low inherent bellman error, International Conference Machine Learning. pp. 10978-10989, PMLR, 2020.
13

---

Published as a conference paper at ICLR 2024
A SUPPLEMENTARY RESULTS FOR SECTION 3.1
A.1 COUNTER-EXAMPLE
LinearQ Definition. LinearQ is class of MDPs, i.e., LinearQ {MM u E N*). Each MDP instance M of LinearQ class is associated with a "size parameter" u. and has state space S" {0, 1.... 3u + 2). action space {0, 1}. and horizon H |S 3u + 3. The initial state is fixed 0.
MD has deterministic transition function T S x A + 8" defined by
s+1
0 S
3u +2
u+ S < 2u and is even
0))
3u+1
u+15 $ < 2u and is odd
3u +
2u+15 5 S 3u + 2,
3u +2
0 S u and is even
3u+1
0 $ S u and s odd
1))
s+1
u+ < 8 I 3u 1 s 3u + 2.
3u +
The reward function r" S" x A [0, 1] is defined by
2k
0 5 $ < u -
1.5k
as
(-2s 28 4u + k)
u u+ S S 2u and s even
28 + 4u + 1.5)
u u+ < 8 < 2u and 8 is odd
2u + 1 - 8 - 3u + 2,
(-8 + 3u + 1.5)k
0 S 8 S u and 8 even
(-8 + 3u + 1)
0 S 8 Σ u and 8 is odd
r as = 1)-
+ 1 < < 3u + 1
s 3u 2.
Here A. Tu+1 is used to scale the reward to range [0, 1].
Properties. The name "LinearQ" comes from its Q-funcioon, which can be represented by a ReLU function (Figure(rrighh and Le
Lemma 1. For any MDP M E LinearQ with size parameter u, the optimal Q-function is
(u'(s,a) = [2k" ReLUU- 2u +1)
a= a a= 1.
1.a Re 3u 1.5)
As a result, the optimal value function is
Qui a= = = 23.2 -- 11
0 s S u ++ S S 3u +
(s)* Qu (8,a - (x k2 (- 3u 1.5)
= 3u 2,
and the optimal policy nh Su + A is
(s) (0 0 05850
(3)
# 1 8 < 3u +2
with optimal return Va.((0) = 2 1).
Proof. We prove by induction on state When 3u + 2, the agent will always be in state 3u + with reward 0, so Q".* (3u + 2, a) = 0, for all a E A.
Now suppose Q" (n, 0) = 2 ReLU 2u 1) and Qq...(n, 1) =) k" ReL- 1.5)
for some integer n in [1, 3u + 2]. Then consider state s = n - 1 E [0, 3u + 1]. For simplicity, we
use s' to denote the next state given state and action If a 0, there are possible cases:
14

---

Published as a conference paper at ICLR 2024
0 1. Then s' = 1 and
(*,(s,0) (s 0) + V ()))= =
8= Then s' = u + 1, and
= (s, - + yu (s') = (s, =
(- 2u +1).
-u+lsss 2u and s even. Then s' 3u 2, and
(s, r" (s. vu (s') = r" (s. + + - 2u 1).
. outtsss 2u and & is odd. Then s' = 3u 1, and
(s, = (s, 0) - vu (s') r" (s, 0) + Qt' (3u + 1,1) = 2k"(-s + 2u + 1).
2u + 15 3u 1. Then = 3u + 2, and
(s, r 0) + V" (8') = r
If a there are possible cases: * 0 s < and is even. Then s' = 3u + 2, and
8 1) = rs (** (8') -0= 30++1. 1.5).
05 u and is odd. Then = 3u + 1, and
(s,ss)==r*(( v (*) = r" (s.1) (3u (3 k- 1.5).
uut 1. Then 5' s u, and
(s, rx V (8') = r r (s. 1) + (,+ k- 3u 1.5).
Therefore, (s = n 1,0) 2k ReLU 1) + 2u + 1) and (s = n 1,1)
K. ReLUR-((((( 1) + 3u 1.5). Thus we complete the proof by induction.
Dataset. The dataset D contains 4(3u + 3)n trajectories, where N* is hyper-parameter for dataset size. 3(3u + 3)n trajectories are sampled by the optimal policy n** We further define (3u + 3) "one-step-deviation" policy (s) (t E S), with each policy sampling trajectories:
(s)
π 11
8 $
8= L.
Simulation Experiment We test the performance of network size, and record the gap between achieved return and optimal return. During evaluation process of MLP-RCSL. we input the exact optimal return desired RTG. We also present the performance of a naive policy as baseline, which always takes action 0 regardless of current state. Experiment details and hyperparameter choices are listed in Appendix E.. As presented in Table 3 MLP-RCSL reaches optimal behavior with constant step size, while Q-learning's best performance only matches the naive policy, even if the network hidden layer size increases linearly with state space size.
and Q-learning with various
A.2 PROOF FOR THEOREM 1
Consider LinearQ (Appendix A.I as our counterexample, then (1) and (ii) hold directly.
For the rest of the theorem, need to first introduce some basic lemmas concerning the relationship of ReLU functions with indicator function in integer domain. The results below are easy to verify and thus omit the proof.
15

---

Published as a conference paper at ICLR 2024
Table 3: Simulation results for LinearQ with different size parameter "Naive" stands for the policy that always takes action 0. "RCSL-w" means MLP-RCSL with hidden layer size w, "QL-w" means Q-learning with hidden layer size w. We report the difference between achieved return on the last training epoch and the optimal return. RCSL achieves optimal return with 16 hidden layers, while Q-learning only achieves return same as naive policy.
Naive
RCSL-16
QL-16
QL-0.5u
QL-u
QL-2u
QL-4u
16
32
32.5
48
64
1
1
1
80
1
1
96
0
1
50.5
112
128
144
126
160
Lemma 2. For any 2, u E L, we have
1x5u = ReLU(-r+ u) + ReLU(-x+uur u+
1 122u ReLU(x u) ReLU u+
1 1z= = ReLUU-= u) + (
-ReLU(x- ReLU(ex-
Proof of (1i2) in Theorem 7 We first observe some critical facts of the dataset D For any state its
sub-optimal action ".* (s) only appears in trajectory sampled by policy n., which achieves RTG Qu, 1 (s)) state since actions after state are all optimal. Thus retum-conditioned policy (s, g) the form
(s) T (s, g) 11-T" n Th)(s))
g Q" (s, 1- n² (s))
== Q (8, --";((
achieves zero training error on D
Define f(x) = 1,20.5, Vor E IR, which the output projection function MLP-RCSL. We show that
x (8,9) g) = f(1g+s=au+1. 1=ts=2u+1 21 21,=0 1sgu + 1)
(4)
via case analysis on state (MLP only needs to represent the input to f):
0 < < Then the optimal action ". (s) 0. We have 1.5. 1. In addition,
have 1 - = O, because the reward on state [0, ai] is always positive and thus the RTG g must be positive.
Now if + == 3u + 1.5, then 1 2=+==uu+1 0, and RHS of (4 is 1, which equals π ( g) as g u. (s. 1). + 3u + 1.5, then RHS of always 0, which also the same
as (s.g).
1 < 2u + Then the optimal action is π **)) = 1. We have 1,5a = 0. In addition, we have 1,=0- = 0, because the reward on state E [0, u] is always positive and thus the RTG g must be positive.
Now if g + 5- = 2u + 1, then g Qu. (s, If
1 +=+==3m 3.. = and RHS of O, which equals (s, g)
3u + 1.5, then RHS of always 1, which is also the same
as (s.g).
2u + 2 < $ < 3u + Then the optimal action x <s)((s) = 1. We have 1.gu = 0. In because 2 g 2u 2.
addition, we have 1 1
Now if g = 0, then RHS of (4) is always O, which equals T(s. g) as = (s, g * 0, then RHS of (4) is always which is also the same =(, g).
16

---

Published as a conference paper at ICLR 2024
Therefore, Equation holds. By Lemma 2 each indicator function in (4) can be represented with ReLU functions with input and g. As a result, we can achieve zero training error with 16 hidden neurons.
We can show that the optimal Q-function can be represented by hidden neurons. Notice that have for all E S and E A = ( that
nd (s, a) 2 ReLU( 2u + 1) 1a=0 kM ReLU(- + 3u 1.5) 1a=1
24 (eLU(-s (2u + 1)aa + 2u + 1) + kn ReLU(-s (3u + 1.5)a).
Thus Qu (s, a) can be represented by 2 hidden neurons, i.e. ReLU(-s (2u + 1)a + 2u + 1) and ReLU(-s (3u 1.5)a).
Proof of (iv) in Theorem 1 [u + 1,2u], we have
We just need to consider the case when a = Notice that for s
(s,a= 0) = Q" (s,a V". (T a (- (-2- 4u + 2)
is even is odd.
1: 4u + 1.5)
For any array (o(k)) (11 l2 2), we recursively define the t-th order difference array of {u(k))kmm, (t S l2 - 41), denoted by { (A)))
12 k=l1" as follows:
{ A v(() ((1-1)(( + 1) A (k)), Vt 2 1,4 sksl2
A (0) o(k) v(k),
1h S 12-
We fix u and construct an array {R(+)) where R(k) = r = ki 0) for E (u,
1. .... 2m). Then its second-order difference array
(2) R()) = (- Vu + 2u 2,
which has t - 2 nonzero terms.
On the other hand, the Q network must be able represent reward function in order to fulfill Bellman completeness, as is discussed in proof idea of Section 3.1 Suppose the Q-network has m hidden neurons. Then it can represent functions the form
m Σ w,ReeU 3aa +
f(s,0)
where a,, B1, c (1 S i S m) as well are all real numbers. Define array { {R(()} k=u+1
and a series of arrays (k)) (kk = Ray cc)) for any k E {u, uu,u+1 1, ... 2u). i E [m]. Then the second-order difference
(1 S I m), in which R(()) f(s = k a 0) and
array of { [R())} k=u+1
2u
m
A (2)
Σ L.),
L,
+1 2u -2
Notice that for any i [m]] there exists most two k's in [u + 1, 1,2u + 2), such that (2)) 0. This indicates that A (2) R(k) has at most 2m non-zero terms.
To make f(s, a) r(s, a) for any s E S", A, we must have (R(k)) (((()) 2u == and thus (2) 22 ==++ (2(2) ((k)) Therefore, we have 2m > u 2, indicating that there least " 2((Sa)) hidden neurons.
17

---

Published as a conference paper at ICLR 2024
A.3 EXPERIMENT
We empirically validate the claim of Theorem on an environment with continuous state and action space. To be specific, We compared RCSL and CQL Point Maze task (Section 5.1 with expert dataset (referred to as "Pointmaze expert"). The expert dataset contains two kinds of trajectories: 1) S - A B and 2) S + M - G (optimal trajectory). Please refer Figure 3 for the maze illustration. The highest return in dataset is 137.9.
Both RCSL policy and CQL Q-network are implemented with three-layer MLP networks, with every hidden layer width equal to paper-prrraeter "dim". We recorded the averaged returns on 4 seeds under different choices of 'dim'.
Table 4:
Comparison of RCSL and CQL on Pointmaze expert under different model size. We recorded the averaged return 4 seeds.
dim=64
dim=128
dim=256
dim=512
dim=1024
RCSL
89.6
97.8
93.3
113
105
CQL.
11.5
11.7
14.7
37.8
23
As shown in Table CQL performance is much worse than RCSL on all model sizes. This is because Bellman completeness is not satisfied for CQL, even with a large model. In contrast, realizability of RCSL can be easily fulfiiled with relatively small model. The experiment supports the theoretical result of Theorem that RCSL can outperform DP-based offline RL methods on deterministic environments, RCSL avoids Bellman completeness requirements.
B PROOFS FOR SECTION 3.2
B. PROOF OF THEOREM
2
Construction of MDPs. Fix any integer constant Hg 2 1. Consider two MDPs M1, M2 with
horizon H 2Hg specified follows:
State space: Both have only one state denoted by s.
Action space: Both have two actions, a, and a
Transition: The transition is trivial as there is only one state.
Reward: There are two possible rewards: good reward rg and bad reward rb, where Tg rb. For M,, and all h E [Ho), we have
12h.M, (s, aL) = rg
M, (s, aR) =
12h 1,M, (5)))) = Tb
12h 1,Ms (8, an) r'g
For M2 we have for all h E [2Ho]]
Th.Ma (S, aL) Tgg
Th an) = Tb-
Construction of Dataset. Fix any integer constant Ko 2 We construct datasets with K 2Koo trajectories for Mi and M2 respectively.
For M1, the trajectories are sampled by two policies: One policy always takes as and the other always takes aR- That is, there are two possible RTG JTG-trajectoriis:
7 = (s, Hor, ar, 8, (Ho 1)r + rgg aL,
8g az)
7 = (s, HoT aR 8, (Ho - 1)F + rb, aR ... Tb, QR).
Let D, the dataset containing Kn 7 and K, r?'s, and DJ be the resulting trajectory-slice dataset with context length for Dj.
18

---

Published as a conference paper at ICLR 2024
For M2,, the trajectories are also sampled by two policies. Both policies alternate between actions ay and ag, one starting with ar and the other starting with ag. Then there are two possible RTG-trajectories
T2 = (s, HoT, aL, 8, (Hl - 1)T + rb. QR. 5, aR)
2 = (s, HoT, an, 8, (Ho 1)r rg. "L," S, a1).
Let D2 be the dataset containing Ko 7 and Ko T²'s, and D2 be the resulting trajectory-slice dataset with context length for D22.
It is obvious that the datasets are compliant with their corresponding MDP. We also point out that DJ = D. other words, the trajectory-slicc datasets with context length for M1 and Ma are exactly the same.
Optimal returns and policies. For both MDPs, the optimal expected RTG g" = Hrg. and the optimal policy is to take action with reward "g at every step.
Analysis of RCSL algorithms Consider Any RCSL algorithm with context length 1. Suppose outputs π an RCSL policy with context length 1, given input DJ- Dj. At step the conditioned RTG is g" for both Mi and M2. Suppose n takes as, with probability and aR with probability
1-p 1-p when conditioning on 9". where 0 S p 1. Then the expected reward at step is prs pr
in MM and pry + (1 p)r.. Since
[pr + (1 p)rgl + (prg + (1 p)r:] = rb + Tg:
there exists i E {1,2), such that the expected reward at step 1 is at most ( + rr) in Mj. Then the expected total return in M; is at most (ro + +rg) = (H 1)rg. so
Ig" J( g-)l 2 2(1, - rr)- Here J=(T',g") ( is the expected return of policy x' conditioning on desired RTG g" in MDP M, In particular, let rig = and rb = 0, then we have |g - Je(m "lmltgg)) 2 2. So we complete the proof.
B.2 RCSL FAILS TO STITCH TRAJECTORIES
In this section, we give a counterexample to show that the decision transformer (Chen al. (021). notable representative of RCSL, fails to stitch trajectories even when the MDP is deterministic. Even when the dataset covers all state-action pairs of this MDP, the decision transformer cannot act optimally with probability 1.
We abuse the notation of trajectory: for any T in the dataset D, we say its sub-trajectory T|i :]] is trajectory in D, written as 7[i j] D.
Here we make an assumption of the generalization ability the return-conditioned policy. When faced with trajectory not in the dataset, the policy outputs mixture of the policy for trajectories in the dataset. For a discrete state space, we assume that the policy only use information related to cach state, not ensembling different states. Thus, the used trajectories for mixture policy should have the same last state as that of the unseen trajectory. To formalize, we have Assumption
Assumption 1, Assume that the return-onditioned policy trained on the dataset 'D has a mapping function J' which
maps an unseen trajectory to {(TI, w1), (72, w2).... where T; E D, the last state of Tú equal to the last state of T, and W1, W..... form probability simplex; maps any T E to {(T, 1)}-
Under Assumption 1 the refurn-conditioned policy proceeds in the following way (Algorithm) 3:
Remark 4. For the weighted part of Assumption [ we view interpolation of several trajectories the dataset. Thus, it is intuitive to view the output policy as an interpolation well. For the "same last state" requirement, this is satisfied in the Markovian RCSL framework. This requirement is also without loss of generality for most generic MDPs, whose states independent.
trajectory
19

---

Published as a conference paper at ICLR 2024
Algorithm 3 Generalization behavior of return-condiiioned policy
1: Input: Retum-condittoned policy S x R (AA) with mapping function fD, desired RTG
Ø1-
2: Observe 81-
3: Set T <- (1.81.
4: for h= 1,2..... H do
5: Sample action ~ E (e W:;
6: Observe reward rh and next state 8h+1.
7: Update 9h+11 t 9h Th and + (T, (T. 1, 8/ Ih+1).
Now we are ready to prove Theorem
Proof of Theorem Construct M in the following way:
Horizon: H 2.
State space: {81, 82, 83).
Action space: A {ai,
Initial state: The initial state is 81, i.e., u(ss) 1.
Transition: For any state sili {1, 2}, any action az(j {1,2)), have ((5)+1();,g)) = 1. For we have P(s3|s3, a,) =
Reward: For any state s:(i E (1,2,3)), any action az(i E (1,23), we have r(8), j
The illustration is shown in Figure 5
T= 0
81
82
83
T=1
r=1
Figure 5: The counterexample showing that decision transformers cannot stitch even in deterministic MDP. Transitions of a are plotted in red arrows, and those of d2 are blue.
The optimal policy to take a2 at both 8 and s2, with total reward of 2.
The istaset
D (81. 1, 013 82. 1, 83, 0). (81, 1, 02; 82, 0, a1 83, 0)}-
All the state-action pairs are covered by D, but the optimal trajectory is not D. Both (3), 1, and (s1, 1, 1, 02) are in D, and they have the weight in the loss function of decision transformer:
(T( log π (a $1, 1) - log (an]s1 1, a2; $2, 0).
log (a, s1, 1) log ( (1 1, 52, 1)
After minimizing the loss function, both a₁ and 02 have non-zero probability at step 1.
When want to achieve the optimal return-to-go 2, the trajectory is (s1,2), 2). The return- conditioned policy æ is fed with a mixture of trajectories in f(r) = {[(s}, 1), 1]]. because r & D. Thus, the policy has non-zero probability mass on the sub-optimal action di.
n
20

---

Published as a conference paper at ICLR 2024
C ABLATION STUDIES
In this section, we study the following questions about MBRCSL:
(1) How does a low-quality dynamics model or behavior policy affect the performance of MBRCSL?
What is the relationshape between MBRCSL performance and the rollout dataset size?
How does the final output policy affect the performance of MBRCSL?
How does data distribution in offline dataset influence the performance of MBRCSL?s
C.I DYNAMICS AND BEHAVIOR POLICY ARCHITECTURE
We tested on Pick and Place task answer question (1). We designed three variants of MBRCSL:
"MIp-Dyn" replaces the transformer dynamics model with an MLP ensemble dynamics model.
"MIp-Beh" replaces the diffusion behavior policy with a MLP behavior policy.
"Trans-Beh" replaces the diffusion behavior policy with a transformer behavior policy.
According to T all three variants fail to reach a positive success rate. The potential reason is that all of them fail generate optimal rollout trajectories. For Myp-Dyn, the low-quality dynamics could mislead behavior policy into wrong actions. For MIp-Beh and Trans-Beh, both of them applies a unimodal behavior policy, which lacks the ability the stitch together segments of different trajectoreis.
Table Ablation study about dynamics and behavior policy architecture on Pick and Place task. We use 4 seeds for each evaluation.
Task
MBRCSL (ours)
Mip-Dyn 0E0
MIp-Beh
Trans-Beh
PickPlace
0.48
0
0
C.2 ROLLOUT DATASET SIZE
To answer question (2), we adjusted the size of rollout dataset, which consists of rollout trajectories with return higher than the maximum return in the offline dataset. As shown in Figure the performance increases significantly when rollout dataset size less than 3000, and then converges between 0.45 and 0.5 afterwards. This is potentially because small rollout datasets are sensitive less accurate dynamics estimate, so that inaccuarate rollouts will lower the output policy's 's performance.
C.3 OUTPUT POLICY ARCHITECTURE
To answer question (3), we investigate the strength of RCSL on Point Maze task by replacing the final RCSL policy with the following architectures:
CQL policy. That is, we train CQL on the rollout dataset, labeled as "MBCQL"
Return-conditooned Gaussian policy (Emmons et al., 22022), labeled as "MBRCSL-Gaussian""
Decision Transformer (Chen et al. 2021). labeled as "MBDT".
ML. policy trained via behavior cloning on the rollout dataset, labeled as "MBBC".
As shown in Table 6 RCSL output policy achieves higher averaged return than CQL with lower variance, demonstrating again the advantage of RCSL over DP-based methods on near-expert datasets. The Gaussian policy captures less optimal behaviors due to the modeling of stochasticity. Decision transformer output policy achieves lower performance than the MLP return-conditooned policy used in MBRCSL. MBRCSL achieves higher expected return than MBBC, because RCSL extracts the policy with highest return from the rollout dataset, while BC can only learn an averaged policy in rollout dataset.
21

---

Published as a conference paper at ICLR 2024
0.50
0.45
0.40
Performance
0.35
0.30
0.25
0.20
0.15
1000
2000
3000
4000
5000
6000
7000
8000
9000
Rollout dataset size
Figure 6: Ablation study about rollout dataset size on Pick and Place task. We use seeds for each evaluation.
Table 6: Comparison between different output policies on Point Maze. All algorithms use rollout dataset collected by MBRCSL We use seeds for evaluations.
Task Pointmaze
MBRCSL (ours)
MBCQL 49.0422.1
MBRCSL-
BBRCSL-Gaussian Gaussian
MBDT
MBBC 83.0
91.5.17.1
62.7431.4
1..
C.4 DATA DISTRIBUTION
Finally, we study the effect of data distribution on Point Maze task. We adjust the ratio between the two kinds of trajectories Point Maze dataset (cf. Figure 3 To be concrete, let
ratio Number of trajectories 5 A
(5)
Offline dataset size
We recorded the averaged performance w.r.t. different ratio value. In addition, we investigated the effect of ratio on rollout process, so we recorded "high return rate", i.e., the ratio between rollout dataset size (number of rollout trajectories with returns larger than the maximum in offline dataset), and the total number of generated rollout trajectories. During the experiments, we fixed the rollout dataset size.
According Figure 7 with ratio. This is potentially because the rollout dataset consists only of rollout trajectories whose returns higher than the maximum offline dataset, i.e., nearly optimal. However, the high return rate increases when ratio approaches 0.5, i.e., uniform distribution. In a skewed dataset, less frequent trajectories will attract less attention the behavior policy, so that the behavior cannot successfully leam those trajectories, leading to drop in the ratio of optimal stitched trajectories in rollouts.
that the averaged return does not show evident correlation
D BELLMAN COMPLETENESS CANNOT BE SATISFIED BY SIMPLY INCREASING MODEL SIZE
We claim that the poor performance of DP-based offline RL baselines in Section is not due to small model size. For this aim, we tested CQL on Point Maze tasks with different model size. The Q-network is implemented with 4-layer MLP network, with each hidden layer width equal to hyper- parameter "dim". The table below presents the averaged return of CQL on 4 seeds under different dim values:
22

---

Published as a conference paper at ICLR 2024
14
80
12
Performance
10
60
rate (%)
8
40
6
4-
20
2
0:
0.5
0.7
0.9
0.3
0.5
0.7
ratio
ratio
Figure 7: Ablation study about data distribution on Point Maze. The left subfigure shows the averaged return for different data distribution, and the right subfigure shows the high return rate w.r.t. different data distribution. We use 4 seeds for each evaluation.
Table Influence model size on performance CQL on Point Maze. We use 4 seeds for each evaluation. The averaged return does not increase with model size.
CQL model size
dim=256
dim=512
dim= 1024
dim=2048
PointMaze averaged return
34.8
52.4
43.7
44.2
We can see from Table 7 that the performance CQL does not increase with a larger model. This because Bellman completeness cannot be satisfied by simply increasing model size. In fact, is a property about the function class on the MDP, so improving CQL would require modifying the architecture satisfy this condition. However, is highly non-trivial design function class that satisfies Bellman completeness.
E EXPERIMENT DETAILS
In this section, we introduce all details in our experiments. Code for reproducing the results is available at atttps://githuu. httpa////weeaaaaaooooo//////////
E.1 DETAILS OF POINT MAZE EXPERIMENT
Environment. We use the Point Maze from Gymnasium-Roboiis (https://robottr) farama org/ wrg//evv envs/ ma catpeppont__aaaaeee maze/), which is a re-implementation of D4RL Maze2D environment (Fu et al. 2020). We choose the dense reward setting, in which the reward is the exponential of negative Euclidean distance between current ball position and goal position. We designed customized maze and associated dataset, as illustrated in Section 5.1 We fix horizon be 200.
Baselines. CQL Kumar 2020) (including CQL part in MBCQL), COMBO Yu al. 2021) and MOPO Yuu et al. 2020) implementations are based on OfflineRL-Kii (https: /ettps//gittut com/ yihaosun 1124 2224/fffiinelli
Kit), open-source offline RL library. We keep the default hyperparameter choices. We implemented DT (Chen et al. 2021)am00 as well as %BC based on official Decision Transformer repository (https: // github..b...... com k2
rmmer...) We set context length to 20 and batch size to 256. We condition initial return of DT the maximum dataset return. Other hyper-pane syper-parameters are the same as Chen et al. (202 Table 9).
MBRCSL. The ensemble dynamics is implemented based on OfflineRL-Kit (nttps: Kit). D. We keep default hyperparameter choice. The dif- fusion behavior policy is adapted from Diffusion Policy repository (https://github com/ real-sttanforrddddfffe lttrifusiioooooicioeee which is modified to a Markovian policy. We keep default hyperparameter choice, except that 10 diffusion steps. Empirically, this leads to good trajectory result with relatively small computational cost. Hyperparameter selections for output policy listed in Table
com/yylhaccsmm11//fffiiii
23

---

Published as a conference paper at ICLR 2024
Table Hyperparameters of output return-coditioned policy in MBRCSL for Point Maze experi- ments
Hyperparameter
Value
Hidden layers
Layer width
1024
Nonlinearity
ReLU
Learning rate
5c-5
Batch size
256
Initial RTG
Maximum return in rollout dataset
Policy Output Deterministic
E.2 DETAILS OF SIMULATED ROBOTIC EXPERIMENTS
Environment. The environments are implemented by Roboverse (https://githul) com/ avisingh599/ [visiigt559//rboveeree)) roboverse). All three tasks (PickPlace, ClosedDrawer, BlowkedDrawer) use the sparse-reward setting, i.e., the reward is received only at the end of an episode, which is I for completing all the required phases, and otherwise. The action is 8-dimensional for all three tasks.
For PickPlace, the initial of position of object is random, but the position of tray, in which we need to put the object fixed. The observation is 17-dimensional, consisting of the object's position (3 dimensions), object's orientation (4 dimensions) and the robot arm's 's state (10 dimensions). The horizon is fixed to 40.
For ClosedDrawer, both the top drawer and bottom drawer start out closed, and the object is placed inside the bottom drawer. The observation is 19-dimensional, consisting of object's position (3 dimensions), object's orientation (4 dimensions), the robot arm's state (10 dimensions), the coordinate of the top drawer (1 dimension) and the x-coordinate of the bottom drawer (1 dimension). The horizon is fixed to 50.
For BlockedDrawer, the top drawer starts out open, while the bottom drawer starts out closed, and the object is placed inside the bottom drawer. The observation space is the same as that of ClosedDrawer, iee, 19-dimensional. The horizon is fixed to 80.
Baselines. CQL (Kumar based on OfflineRL-Ki Kit (https: rithub..b....
2020) (including CQL part MBCQL), COMBO Yu et 202 are
Kit). We
keep the default hyperparameter choices, except that we set conservative penalty in both CQL and COMBO to 1.0, which leads to higher successful rate for CQL. We implemented DT (Chen al..|22000 (222)) based on official Decision Transformer repository (https: / /ttgt/tttbb...////// com k
:decssion rann formert We set context length to 20 and batch size to 256. We condition initial return of DT on the maximum dataset return. Other hyper- parameters are the same as Chen al. 2021 Table 9).
MBRCSL. The diffusion behavior policy is the same as in Point Maze Experiment, except that use 5 diffusion steps, which has the best performance empirically. The transformer model in wwt.://ococaaaaaaaaaecaeeaee http:///g.......coo//aattttttttt TGFT), an open-source We adapt it into
dynamics is based on MinGPT pottps://githin http:/ub.ccoomppppaaaaa/a// simple re-implementation of GPT dynamics model (s', a). Output policy implemented with autoregressive model. The detailed hyperparameter choices are listed in Table 9
E.3 DETAILS OF LINEARQ EXPERIMENT
The algorithms are implemented with minimal extra tricks, so that they resemble their theoretical versions discussed in Section 3.1
The policy in MLP-RCSL algorithm is implemented with MLP model. It takes the current state and return input, and outputs a predicted action. We train the policy to minimize the MSE error of action (cf. Eq. (2).
The Q-function in Q-learning algorithm is implemented with a MLP model. takes current state s as input and outputs (Q(s, = 0), ((, s, a = 1)), the estimated optimal Q functions on both a = 0 and 1. During evaluation, the action with higher estimated Q-function value is taken. We maintain
24

---

Published as a conference paper at ICLR 2024
Table 9: Hyperparameters of MBRCSL for Fick-and-Place experiments
Hyperparameter
Value
Horizon
40) PickPlace
50 Drawer
80 BlockedDrawer
Dynamice/architecture
transformer
Dynamics/number of layers
4
Dynamics/number attention heads
Dynamics/embedding dimension
32
Dynamics/batch size
256
Dynamics/lcarning rate
lc-3
recaavior/rrciitecurr
diffuser
Behavior/diffusion steps
Behavior/batch size
256
Behavior/epochs
PickPlace,
40 ClosedDrawer, BlockedDrawer
RCSL/architecture
autoregressive
RCSL/hidden layers
4
RCSL/layer width
200
RCSL/nonlinearity
LeakyReIUU
RCSL/lcarning rate
le-3
RCSL/batch size
256
RCSL/Initial RTG
Maximum return in rollout dataset
separate target Qentwwork and update target network every 10 epochs, so that the loss can converge before target network update.
We train both algorithms for 300 epochs and record their performance on the last epoch.
25


Reviews:
[{"review_id": "ToHicm39XK", "reviewer": "Reviewer_YXKR", "summary": "This paper aims to investigate how to make the RCSL method work in the off-policy RL setting. The paper first shows that RCSL can achieve better performance than DP-based algorithms in near deterministic environments, as long as the dataset contain optimal trajectories, and points out that this is due to that there is no need of Bellman-completeness for RCSL. This indicates that this type of methods should work if sufficient data coverage is obtained. Hence the authors propose a model-based RCSL (MBRCSL) algorithm which is able stitch together sub-optimal data while avoiding the Bellman-completeness requirements. Empriical results on D4RL benchmark also verify its effectiveness.", "questions": "1. Theorem 1 only shows that there exists some MDPs that don't need the Bellman-completeness assumption for RCSL to work, but how to characterize such MDPs? How about other environments that lacks such characteristic?\n\n2. In section.4, it's said that 'in many cases, this generated data provides coverge over potentially optimal trajectory and return data....', in my opinion, the policy learned by behavior cloning from suboptimal offline trajectories is highly likely to be a suboptimal policy as well, even though the dynamic model is leanred accurately, would you please make it more clear that why the trajectories sampled by the behavior policy with approximat dynamic model will become optimal?  In particular, even if the training rollouts consist of those synthetic trajectories with a return greater than the ones in offline dataset, these better rollouts are based on the approximate dynamic model and cannot guarantee the optimal trajectory. If so, how can we learn the optimal policy using the RCSL algorithm?\n\n3. Has the authors tried MB+DT or MB+BC, i.e., replacing the RCSL component in MBRCSL with other methods without the requirement of Bellman-completeness?\n\n4. on p.4., below section 3.2, the paper says that \"RCSL cannot perform trajecatory stitching\", would you please make it more clear? In my opinion, RCSL works by supervised learning, in which  part of its data features (future return or goals for each state) are generated by trajactory stitching by hand.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "While previous method work (Brandfonbrener et al., 2022) gives negative theoretical results for RCSL, this paper investigates how to make it work in practice in despite of those disadvantages, while preserving its advantages. The paper is well written, easy to read and demonstrates that the proposed method may significantly outperforms SOTA offline RL methods under certain environments.", "weaknesses": "The proposal to address the 'trajectory stitching' issue seems idealistic - how can you ensure that sufficient data coverage can be obtained simply by rolling out a learned dynamics model?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dbsxyogASI", "reviewer": "Reviewer_HXBB", "summary": "The paper tackles a central challenge in offline reinforcement learning - the difficulty of stitching together suboptimal trajectories into optimal behavior, while using function approximation. Standard off-policy algorithms like Q-learning rely on dynamic programming and the strong requirement of Bellman completeness, which can be problematic to satisfy with function approximation. \n\nIn contrast, the paper shows return-conditioned supervised learning (RCSL) methods that directly train policies on trajectories need only the easier-to-satisfy realizability condition. However, RCSL methods cannot stitch trajectories. To get the best of both worlds, the paper proposes model-based RCSL (MBRCSL), creatively combining RCSL and model-based RL. MBRCSL trains a dynamics model, uses it to sample high-performing trajectories, and trains an RCSL policy on these optimal trajectories. This enables trajectory stitching without requiring Bellman completeness. \n\nTheoretically, the paper sharply characterizes when RCSL can outperform Q-learning in simple settings and why RCSL fails at stitching. Empirically, MBRCSL demonstrates clear improvements over model-free and model-based baselines on Point Maze and Pick-and-Place environments.", "questions": "How does MBRCSL compare to other state-of-the-art offline RL algorithm on benchmark tasks?\n\nHow does MBRCSL compare to prior model-based offline RL methods like MOReL, MOPO, COMBO including sample complexity and compute? E.g. your approach requires accurate learned dynamics model and uncertainty modeling is less sophisticated than ensemble approaches like MOPO. No explicit handling of OOD actions unlike conservative methods means the algorithm could be brittle on highly suboptimal datasets- could you comment on this?\n\nCan you formally characterize conditions under which model rollouts provide sufficient coverage for optimality of RCSL?\n\nCan you incorporate offline constraints into RCSL policy learning?\n\nCan you analyze sample efficiency of MBRCSL with limited suboptimal data?\n\nHow many trajectories do you need before performing BC? I would imagine the performance of the algorithms depends a lot on this?\n\nIn general, I'm quite surprised the algorithms works as explained because if BC is a policy that combines the best and worst trajectories in the dataset I don't understand how it can find trajectories that are better than the best ones. Unfortunately the experiments aren't particularly convincing in this respect.\n\n\n\n\nIs end-to-end training of model and RCSL policy possible?\n\nWhat are the most important architectural choices for dynamics model and RCSL policy?\n\nCan MBRCSL be extended to stochastic environments?\n\nThe theoretical results comparing RCSL and Q-learning are derived in the tabular setting. However, function approximation is crucial for offline RL in complex environments with large state spaces. Can you discuss how the theoretical insights around Bellman completeness and realizability could be extended to the function approximation setting? Analyzing the hidden layer width requirements for MLP-RCSL versus MLP Q-learning could provide useful insights into when RCSL methods might have advantages over dynamic programming with function approximation for offline RL.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Provides concrete theoretical characterization showing RCSL only requires realizability while Q-learning requires stronger Bellman completeness.\n\nProves fundamental limitation that Markovian RCSL policies cannot perform trajectory stitching.\n\nInteresting idea of using model rollouts to generate near-optimal trajectories then training RCSL policy.\n\nModular framework allows flexible choice of model architecture. \n\nGood empirical results on Point Maze and Pick-and-Place. Outperforms model-free baselines.", "weaknesses": "Limited experimental validation on only 2 environments. More complex domains needed.\n\nNo comparison to other state-of-the-art model-based offline RL algorithms\n\nTheoretical results consider tabular case. Unclear if insights extend directly to function approximation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5wgUUTrlSb", "reviewer": "Reviewer_9pmM", "summary": "This paper presents MBRCSL, a return-conditioned supervised learning algorithm that improves on one of the major limitations of that algorithm category, the inability to perform trajectory stitching, by adding model-based rollouts that synthesize multiple trajectories to augment the training data. The authors also prove that under certain assumptions RCSL algorithms converge with weaker assumptions about the contents of the offline dataset compared to offline Q-learning methods.", "questions": "-For the two tasks tested on, the poor performance of CQL and co. is claimed to be due to issues with Bellman completeness. Could it be reasoned about or experimentally demonstrated what additional data is needed to bring those algorithms up to performance levels similar to MBRCSL? I imagine the additions needed are non-trivial, which would be a mark in MBRCSL's favor, but I'm curious to know how hard it is to satisfy the requirements of offline Q-learning algorithms in comparison. Basically, how hard of a constraint to satisfy is the need for Bellman completeness in practice versus theory?\n\n-I'd be interested to see a non-tabular environment testing how model size/parameter count is affected by RCSL versus Q-learning. The theoretical argument here seems strong, so I'm curious if performance as a function of model size ends up being meaningfully different between RCSL methods and Q-learning methods.\n\n-MBRCSL relies on a good forward dynamics model and a behavior policy to allow trajectory stitching, and while this is a fair requirement I'm curious what happens when those two components are lower in quality (for example, a weaker dynamics model versus a stronger one). Does performance degrade smoothly, or is a high-quality model essential?\n\n-In the same vein, what happens if the rollout budget is reduced? How does MBRCSL degrade with limited synthetic data?\n\n-What happens to MBRCSL if dataset coverage is non-uniform? I can imagine something like a constructed dataset that has a tunable degree of biased-ness in coverage being informative here.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I found this paper very interesting- it combines non-trivial theoretical and experimental arguments in a way I don't see often, and does so in a very comprehensible way, speaking as someone who hasn't worked on offline RL and isn't deeply familiar with the literature on the topic.\n\nIn general I liked the ideas presented- they rely on some assumptions about the dataset and MDP(s) of interest, but these are mostly stated up front and seem if not trivial in practice at least a reasonable starting point for future work to relax to some degree. The point about realizability versus Bellman completeness seems well put to me.", "weaknesses": "My biggest criticism of this paper is that it doesn't really have any ablations or demonstrate the experimental limitations of MBRCSL at all. The theoretical assumptions are all clearly stated, but it's not clear how far these will stretch in practice.\n\nI think what's already here is clearly argued and significant, so I'm inclined to recommend acceptance regardless, but some more stress tests to see where assumptions break down would greatly improve the paper. I've added some questions below in that vein, hopefully they can be useful for either this paper or future work building on it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XiWpkh0nyt", "reviewer": "Reviewer_BsaB", "summary": "The paper proposes a novel method for off-line reinforcement learning based on return-conditioned supervising learning (RCSL) algorithms for policy selection working on an augmented data set of roll-out trajectories generated by means of a dynamics model learned from actually experienced trajectories. An empirical evaluation demonstrates the advantages of the proposed method on several decision problems. In addition, the paper provided theoretical analysis of why return-conditioned algorithms cannot stitch an optimal policy from segments of trajectories contained in a data set, thus motivating the idea to augment such a data set with rollouts.", "questions": "How exactly is the policy represented in the comparison in Section 3.1? To begin with, there seems to be a discrepancy between the kind of policies introduced in Section 2, where a policy maps trajectories to probability distributions over the action set, and the one in Section 3.1, where the policy is deterministic, and outputs a single action. Equation 2 is compatible only with the latter type of policy, and only if the actions are real-valued. But, the projection f(a) seems to be a thresholding operation, returning only possible values of 0 and 1. They could be used as real-valued actions, but then, how are problems with more than two actions represented? Also, if f(a) is a hard thresholding, it is not differentiable, how is the minimization problem in the MLP-RCSL solved?\n\nAlso, are the indices for tau[i, j] in the first equation in Section 2.2 correct? For cx=1, we get i>j, which seems to contradict the definition of a sub-sequence tau[i, j]. Also, isn't s_h already part of that sequence, why is it listed as a separate argument of the policy?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main strength of the paper is the combination of RCSL with rollouts using a learned model for the purposes of augmentation of a collected data set of trajectories. The use of RCSL allows the use of machine learning models that are only adequate for representing the optimal policy, and not necessarily for representing the optimal Q-function of the problem, too, as required by methods based on dynamic programming (Bellman back-ups). This makes the selection of the neural network models for a given decision problem easier. On the other hand, the use of rollouts produces by means of a learned model increases the set of trajectories that the RCSL will be selecting from, thus increasing the cumulative reward of the resulting policy.", "weaknesses": "The inability of RCSL to stitch an optimal trajectory from multiple segments of trajectories in the data set it is working with seems like a major limitation. If it is true that RCSL can learn the optimal policy only if the optimal trajectory is in the data set, then for optimality, the augmented data set must be large enough to contain this trajectory. But, it is not clear how many rollouts will be needed for this, even in some probabilistic sense.\n\nAnother weakness of the paper is that it is hard to follow the proposed methods. For example, the RCSL method can be explained better.  \n\nA minor typo, in the abstract: \"unde\" -> \"under\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)