Free from Bellman Completeness: Trajectory Stitching via Model-based Return-conditioned Supervised Learning

Thank you for your constructive feedback. We address your comments below.

Weakness 1. 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.

We tested 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. The result is shown in the table below:

The results are averaged on 4 seeds. We can see that the probability is acceptable in our benchmarks considered. It is true that the probability will reduce if the optimal trajectory consists of more trajectory segments in the offline dataset. We will work on this problem as future work.

Weakness 2. It is hard to follow the proposed methods. For example, the RCSL method can be explained better.

MBRCSL works as follows: First, we learn an approximate dynamics model as well an expressive behavior policy from the offline dataset. Then we rollout the behavior policy on the dynamics model to generate synthetic trajectories, which stitch together trajectory segments in the offline dataset. In this process, potentially optimal trajectories will be generated, although the averaged return of behavior policy could be low. We aggregate potentially optimal trajectories into a new rollout dataset by picking out trajectories with estimated returns larger than the maximum return in the dataset. Finally, we apply RCSL on the rollout dataset, to further extract the optimal policy.

RCSL learns a distribution of actions in each state conditioned on a desired return-to-go by directly applying supervised learning on the trajectories in the dataset. RCSL is a standard method studied in [1, 2]. We formalize RCSL framework in Section 2.2 for rigorous theoretical analysis.

[1] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 35:1542–1553, 2022. [2] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning?, 2022.

Question 1. How exactly is the policy represented in the comparison in Section 3.1?

The policy defined in Section 2. is the most general version, which is a stochastic policy. In Section 3.1, we focus on deterministic policies for RCSL. This is without loss of generality, as deterministic policy class includes the optimal policy. Equation (2) is the training objective for MLP-RCSL, which is indeed compatible to deterministic policies only.

In Section 3.1, the MLP-RCSL architecture is exactly the same as described in Section 2.2, and it is trained by minimizing Equation (2). The projection $f(a)$ is a component external to MLP-RCSL. It is added to MLP-RCSL to ensure that the output action lies exactly in the action space ($\{0,1\}$) of the example considered in Theorem 1 (cf. Figure 2). For other MDPs, the projection can be omitted or adapted to the action space. We have revised the formulation in Section 3.1.

Question 2. Are the indices for $\tau$[i: j] in the first equation in Section 2.2 correct?

The indices are correct according to our definition in Section 2, paragraph "Trajectories", in which $\tau[i:j] = (s_i,g_i,a_i,\cdots,s_j,g_j,a_j)$. When cx=1, the term $\tau[h-cx+1:h-1 ]$ in $\pi^{\text{cx}}(a_h|\tau[h-\text{cx}+1: h-1],s_h,g_h)$ is omitted and the resulting policy is $\pi^1 (a_h|s_h,g_h)$. We have added a clarification in Section 2.2.

Weakness 3. A minor typo in the abstract.

Thanks for pointing out. We have fixed it in our paper.

Question 4. What happens if the rollout budget is reduced?

We added an ablation study in our paper, please refer to Appendix C.2. As shown in the table below, we adjusted rollout dataset size (numbers in the first row) and recorded the averaged success rate (numbers in the second row) on PickPlace task.

We can see that the performance increases significantly when rollout dataset size is less than 3000. One reason is that small rollout datasets are sensitive to less accurate dynamics estimate, so that inaccuarate rollouts will lower the output policy’s performance. The performance then converges between 0.45 and 0.52 when rollout dataset size is large than 3000.

Question 5. What happens to MBRCSL if dataset coverage is non-uniform?

We tested the effect of data distribution on Point Maze. We fixed the offline dataset size and adjusted the ratio between the two kinds of trajectories in Point Maze dataset (cf. Figure 3). To be concrete, let $\text{ratio} := \text{Number of trajectories S → A → B → G} / \text{ 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. The results are shown below (All results are averaged on 4 seeds):

We can see that the averaged return does not show an evident correlation with ratio. This is potentially because the rollout dataset consists only of rollout trajectories whose returns are 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 to the behavior policy, so that the behavior cannot successfully learn those trajectories, leading to a drop in the ratio of optimal stitched trajectories in rollouts.

We added these results as an ablation study in our paper, please refer to Appendix C.4.

Thank you for your positive feedback and for finding our paper “interesting” and "comprehensive." We address your comments below.

Question 1. 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?

We tested CQL on Point Maze tasks with different model size. The Q-network is implemented with a 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:

We can see that the performance of CQL does not increase with a larger model. This is because Bellman completeness cannot be satisfied by simply increasing model size. In fact, it is a property about the function class on the MDP, so improving CQL would require modifying the architecture to satisfy this condition. However, it is highly non-trivial to design a function class that satisfies Bellman completeness. We added this discussion in Appendix D of our paper.

Question 2. A non-tabular environment testing how model size/parameter count is affected by RCSL versus Q-learning

We compared RCSL and CQL on Point Maze task with expert dataset (referred to as "Pointmaze expert"). To be specific, the dataset contains two kinds of trajectories: $S\to A \to B$ and $S\to M \to G$ (optimal trajectory), and the highest return in dataset is 137.9. Please refer to Figure 3 of Section 5 for the maze illustration. Both RCSL policy and CQL Q-network are implemented with three-layer MLP networks, with every hidden layer width equal to hyper-parameter dim. We recorded the averaged returns on 4 seeds under different choices of dim:

CQL performance is much worse than RCSL on all model size, though both algorithms tend to achieve a higher averaged return when model size increases (especially for dim$\leq 512$). This result supports the claim in Section 3.1 that RCSL can outperform DP-based offline RL methods, as RCSL avoids Bellman completeness requirements. We added this experiment in Appendix A.3 of the paper.

Question 3. 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?

We have added an ablation study to study the dynamics model and behavior policy, please refer to Appendix C.1. We designed three variants of MBRCSL:

• “Mlp-Dyn” replaces the transformer dynamics model with an MLP ensemble dynamics model.
• “Mlp-Beh” replaces the diffusion behavior policy with a MLP behavior policy.
• “Trans-Beh” replaces the diffusion behavior policy with a transformer behavior policy.

We tested on PickPlace environment and got the following results:

We can see that both the quality of dynamics and that of behavior policy are crucial to final performance. One reason for failure of the three variants is that none of them are able to generate optimal trajectories. For Mlp-Dyn, the low-quality dynamics will mislead behavior policy into wrong actions. For Mlp-Beh and Trans-Beh, both of them apply the unimodal behavior policy, which lacks the ability the stitch together segments of different trajectories.

Thanks for the additional explanation and experimental validation! These additional tests give me more confidence in the experimental results, and as such I've raised my score. I think with these additions the paper is in a good place now.

Question 8. Is end-to-end training of model and RCSL policy possible?

An end-to-end training scheme is not feasible in our setting. As illustrated in Figure 1, the training order is: dynamics model & behavior policy $\to$ trajectory rollout $\to$ return-conditioned output policy. This is because training of return-conditioned output policy depends on rollout data, which is generated through dynamics model and behavior policy.

Question 9. What are the most important architectural choices for dynamics model and RCSL policy?

The dynamics model must be accurate enough, so we recommend using a transformer model, especially in complex tasks such as the simulated robotic tasks in Section 5.2. On the other hand, the RCSL policy architecture could be task dependent. Please refer to [1] for details.

We've also added an ablation study to investigate different architectures of dynamics and RCSL policy in our paper. Please refer to Appendix C.1 and C.3.

[1] Yang, S., Nachum, O., Du, Y., Wei, J., Abbeel, P., & Schuurmans, D. (2023). Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.04129.

Question 10. Can MBRCSL be extended to stochastic environments?

As written in our conclusion, we will consider MBRCSL in stochastic environments as a direct future work. One possible solution is to learn a more general trajectory representation other than return-to-go, as described in [1, 2].

[1] Paster, Keiran, Sheila McIlraith, and Jimmy Ba. "You can’t count on luck: Why decision transformers and rvs fail in stochastic environments." Advances in Neural Information Processing Systems* 35 (2022): 38966-38979.

[2] Yang, M., Schuurmans, D., Abbeel, P., & Nachum, O. (2022). Dichotomy of control: Separating what you can control from what you cannot. Preprint arXiv:2210.13435.

Question 3. Can you formally characterize conditions under which model rollouts provide sufficient coverage for optimality of RCSL?

If the dynamics is accurate enough and the behavior policy is multi-modal, then model rollouts would be able to generate optimal stitched trajectories with certain probability. Specifically, if we obtain a uniform convergence of our estimated model for policy evaluation in the reward-free setting (which guarantees for any reward function, we can obtain an optimal policy) [1]. Then we can choose rewrds to be the indicator functions for each state, and the induced policies will give us sufficient coverage.

[1] Ming Yin, Yu-Xiang Wang. Optimal Uniform OPE and Model-based Offline Reinforcement Learning in Time-Homogeneous, Reward-Free and Task-Agnostic Settings. Advances in Neural Information Processing Systems, 2021.

Question 4. Can you incorporate offline constraints into RCSL policy learning?

RCSL does not need explicit offline RL constraints like most dynamic programming based offline RL methods like CQL, COMBO and MOPO. This is because RCSL is inherently weighted supervised learning, which does not need to handle OOD actions.

Question 5. Can you analyze sample efficiency of MBRCSL with limited suboptimal data?

Thanks for asking. We believe it is possible to combine the analysis of model-based offline RL [1] and recent development for RCSL [2]. We leave it as a future work to give a formal sample complexity theorem of MBRCSL.

[1] Ming Yin, Yu-Xiang Wang. Optimal Uniform OPE and Model-based Offline Reinforcement Learning in Time-Homogeneous, Reward-Free and Task-Agnostic Settings. Advances in Neural Information Processing Systems, 2021.

[2] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 022

Question 6. How many trajectories do you need before performing BC? I would imagine the performance of the algorithms depends a lot on this?

We added an ablation study on PickPlace task, as shown in the table below. We adjusted rollout dataset size (numbers in the first row) and recorded the averaged success rate (numbers in the second row).

We can see that the performance increases significantly when rollout dataset size is less than 3000, and then converges between 0.45 and 0.52 afterwards. We have added this ablation study to our paper in Appendix C.

Question 7. In 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.

The behavior policy does not discern the quality of generated rollout trajectories. Instead, it only needs to stitch trajectory segments in offline dataset. As an illustration, suppose an MDP has 5 states $\mathcal S=\{A,B,C,D,E\}$, and the dataset consists of two kinds of trajectories: $A\to C \to E$ and $B\to C \to D$. Then by rolling out the behavior policy, we obtain 4 possible kinds of rollout trajectories: $A\to C \to E$, $A\to C \to D$, $B\to C \to D$, $B \to C \to E$. Some of these trajectories have high returns, and others have low returns. Given heterogeneous rollout trajectories, the output policy is able to extract the best trajectories due to return-conditioning [1-3]. To summarize, the behavior policy only provides stitched rollout trajectories with varied performance, and the return-conditioned output policy extracts the best policy from rollout dataset.

[1] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 35:1542–1553, 2022.

[2] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning?, 2022.

[3] 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:15084–15097, 2021.

Thank you for your careful review and constructive feedback. We address your questions below.

Weakness 1. Limited experimental validation on only 2 environments.

We added two simulated robotics tasks, "ClosedDrawer" and "BlockedDrawer". The former requires the robot to accomplish (1) Open a drawer; (2) Grasp the object in the drawer. The latter requries the robot to do (1) Close a drawer on top of the target drawer, and then open the target drawer; (2) Grasp the object in the target drawer.

The results are shown below. We also added the result to our paper in Table 2 in Section 5.

Weakness 2. No comparison to other state-of-the-art model-based offline RL algorithms.

The original submission featured comparisons to a few state-of-the-art model-based offline RL baselines (COMBO & MOPO for Point Maze, COMBO for PickPlace.) We additionally compare to MOReL (Kidambi et al., 2020) on Point Maze in Table 1 of our revised submission, and found MBRCSL to outperform MOReL by a significant margin.

Weakness 3. Theoretical results consider tabular case. Unclear if insights extend directly to function approximation. Question 11. The theoretical results comparing RCSL and Q-learning are derived in the tabular setting.

Although the class of MDPs used in Theorem 1 have finite state space, we indeed considered function approximation in our analysis. In particular, the realizability and Bellman completeness properties are discussed in the function approximation setting. Our theoretical analyses in Section 3 were also carried out under function approximation. In Theorem 1, the RCSL policy adopts an MLP network architecture instead of tabular representation, and $Q$-Learning also represents the $Q$-function with an MLP network.

In Theorem 2, we laid no assumptions on the RCSL policy architecture, i.e., it can take any form of function approximation. In Theorem 3, we consider Decision Transformer as RCSL policy architecture.

In addition, the example used in Theorem 1 can be adapted to an MDP with continuous state/action space by adding a continuous dummy variable. In general, consider an MDP $\mathcal M = (H, \mathcal S, \mathcal A, \mu, \mathcal T , r)$, where state space $\mathcal S$ and action space $\mathcal A$ are all discrete . We can construct a new MDP $\mathcal M' = (H, \mathcal S', \mathcal A', \mu', \mathcal T', r')$ with state space $\mathcal S' = \mathcal S \times \mathbb [0,1]$ and action space $\mathcal A' = \mathcal A \times \mathbb [0,1]$. The initial state distribution PDF $\mu'(s,x) = \mu(s)$, for any $(s,x)\in \mathcal S'$. The transition dynamics is $\mathcal T'(s',x'| s,x,a,y) = \mathcal T(s'|s,a)$, for any $(s',x'), (s,x)\in \mathcal S', (a,y)\in \mathcal A'$. The reward function is $r'(s,x,a,y) = r(s,a)$, for any $(s,x)\in \mathcal S', (a,y)\in \mathcal A'$. In fact, $\mathcal M'$ is equivalent to $\mathcal M$, so Theorem 1 holds on $\mathcal M'$ if and only if it holds on $\mathcal M$. This means that Theorem 1 also applies to MDPs with continuous state and action space.

Question 1. How does MBRCSL compare to other state-of-the-art offline RL algorithm on benchmark tasks?

In Tables 1 and 2 of our original submission, we compared MBRCSL with state-of-the-art offline RL algorithms, including model-based approaches e.g. COMBO, model-free approaches e.g. CQL and RCSL methods e.g. DT. We additionally compare to MOReL on D4RL pointmaze in our revised submission. We found MBRCSL to outperform all of the baselines in our experiments.

Question 2. How 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?

Offline RL methods based on dynamic programming approaches, e.g., MOReL, MOPO and COMBO requires conservatism or pessimism to avoid over-estimation of $Q$-functions for OOD actions. In contrast, MBRCSL does not rely on dynamic programming and thus does not need to handle OOD actions. In fact, the behavior policy trained via supervised learning will only take in-distribution actions, so the dynamics model can give a good estimation of states and rewards without additional conservatism. According to Table 1 of Section 5, MBRCSL indeed outperforms DP-based offline RL methods empirically.

Dear Reviewer,

Thank you for your time and efforts in reviewing our work. We have provided detailed clarification to address the issues raised in your comments. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.

If you have any additional questions or comments, we would be happy to have further discussions.

Thanks,

The authors

Question 4. 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 trajectory stitching by hand.

The return-to-gos in standard RCSL are Monte-Carlo, in that they are generated from a single trajectory in the offline dataset. On the other hand, trajctory stitching requires combining multiple trajectories. During training, the return-to-go of a trajectory $\tau$ at timestep $h$ is computed by accumulating the rewards achieved at timesteps $h'\geq h$ (cf. Section 2, paragraph "Trajectories"). During evaluation (Algorithm 1), the desired return-to-go $\tilde g_1$ can be assigned either as the highest return in dataset, or a pre-specified return to achieve according to the task. Our definition for the training and evaluation processes of RCSL aligns with previous RCSL literature [1-3].

[1] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 35:1542–1553, 2022.

[2] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning?, 2022.

[3] 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:15084–15097, 2021.

Thank you for your review. We address your questions below.

Question 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?

For any MDP, Bellman completeness is a property of the function approximation class (cf. Definition 2). For any MDP, the tabular function class always satisfies Bellman completeness, as it encompasses all possible $Q$-functions. It is easier for a function class to satisfy the realizablity requirement of RCSL than the Bellman completeness requirement of $Q$-learning because the realizabity approximation error ($\max_{Q \in \mathcal{F}}\|Q-Q^*\|$) decreases monotonically as the function class size increases, but this is not true for Bellman completeness approximation error. Theorem 1 proves this belief rigorously with a concrete example.

Question 2. Would you please make it more clear that why the trajectories sampled by the behavior policy with approximat dynamic model will become optimal? Weakness. How can you ensure that sufficient data coverage can be obtained simply by rolling out a learned dynamics model?

Although the behavior policy learned from offline dataset is sub-optimal in terms of expected return, it is able to stitch trajectory segments in offline dataset. As an illustration, suppose an MDP has 5 states $\mathcal S=\{A,B,C,D,E\}$. The dataset consists of two kinds of trajectories: $A\to C \to E$ and $B\to C \to D$, and the optimal trajectory is $A\to C\to D$. Then there will be 4 possible kinds of rollout trajectories: $A\to C \to E$, $A\to C \to D$, $B\to C \to D$, $B \to C \to E$. Therefore, the behavior policy is able to generate the optimal trajectory, albeit with low probability. The return-conditioned output policy to find out the best trajectories in rollout dataset, by applying return-conditioned supervised learning.

In fact, we tested 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, as shown in the table below:

We can see that i.i.d. rollout through dynamics model indeed generates adequate high-return trajectories empirically.

While the returns of rollout trajectories are generated by the learned dynamics model, the estimation error can be reduced to an acceptable rate by learning an accurate dynamics. We provide an ablation study on dynamics model architecture in Appendix C.1.

Question 3. 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?

We added the two experiments on Point Maze environment, as shown in the table below. We have updated this result in our paper, please refer to Appendix C.3.

We can see that Decision Transformer output policy achieves lower performance than the MLP return-conditioned 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.

Dear Reviewer,

Thank you for your time and efforts in reviewing our work. We have provided detailed clarification to address the issues raised in your comments. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.

If you have any additional questions or comments, we would be happy to have further discussions.

Thanks,

The authors

If my understanding is correct, what this paper does is to propose a preprocessing steps for RCSL, which is essentially a data augumentation prodedure based on a learnt model. In this way, the problem that "RCSL CANNOT PERFORM TRAJECTORY-STITCHING" could potentially be addressed. But I'm still curious about why other offline methods, e.g., Decision Transformer cannot benefit from those roll outs. In other words, I think that the close connection between the proposed method and RCSL is not well revealed in this paper although it shows that such combination indeed yields good empirical results. Besides, employing rollouts from a model is a popular way for model-based offline RL, and in this sense, the novelty of the proposed method could be further clarified.