Paper: Free from Bellman Completeness: Trajectory Stitching via Model-based Return-conditioned Supervised Learning (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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:

| Tasks | Pointmaze | PickPlace | ClosedDrawer | BlockedDrawer |
| ----- | --------- | --------- | ------------ | ------------- |
| ratio | 14.3%     | 19.2%     | 25.7%        | 63.2%         |

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.

Reviewer comment (Writing):
Another weakness of the paper is that it is hard to follow the proposed methods. For example, the RCSL method can be explained better.

Author response:
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.

Reviewer comment (Writing):
A minor typo, in the abstract: "unde" -> "under"

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

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
> **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:

| Task      | dim=256 | dim=512 | dim=1024 | dim=2048 |
| --------- | ------- | ------- | -------- | -------- |
| PointMaze | 34.8    | 52.4    | 43.7     | 44.2     |

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.

Reviewer comment (Experiments):
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.

Author response:
> **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`:


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

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.

Reviewer comment (Experiments):
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?

Author response:
> **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:

| Task      | MBRCSL (ours) | Mlp-Dyn | Mlp-Beh | Trans-Beh |
| --------- | ------------- | ------- | ------- | --------- |
| PickPlace | 0.48±0.04     | 0 ± 0   | 0 ± 0   | 0 ± 0     |

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.

Reviewer comment (Experiments):
What happens if the rollout budget is reduced? How does MBRCSL degrade with limited synthetic data?

Author response:
> **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. 

| Task      | 1000   | 2000   | 3000   | 4000   | 5000   | 6000   | 7000   | 8000   | 9000   |
| --------- | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ |
| PickPlace | 0.1650 | 0.3875 | 0.4775 | 0.4525 | 0.4775 | 0.4475 | 0.4950 | 0.5125 | 0.4825 |

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.

Reviewer comment (Experiments):
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.

Author response:
>  **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):

| ratio            | 0.1   | 0.3    | 0.5    | 0.7   | 0.9   |
| ---------------- | ----- | ------ | ------ | ----- | ----- |
| return           | 88.9  | 79.9   | 91.5   | 77.2  | 82.9  |
| high return rate | 0.29% | 13.94% | 14.25% | 1.30% | 0.47% |

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.

Reviewer comment (Experiments):
Limited experimental validation on only 2 environments. More complex domains needed.

Author response:
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.

| Task          | MBRCSL (ours) | CQL       | COMBO | DT   | BC        |
| ------------- | ------------- | --------- | ----- | ---- | --------- |
| ClosedDrawer  | **0.51±0.12** | 0.11±0.08 | 0±0   | 0±0  | 0.38±0.02 |
| BlockedDrawer | **0.68±0.09** | 0.34±0.23 | 0±0   | 0±0  | 0.61±0.02 |

Reviewer comment (Experiments):
No comparison to other state-of-the-art model-based offline RL algorithms.

Author response:
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.

Reviewer comment (Theory):
Theoretical results consider tabular case. Unclear if insights extend directly to function approximation.

Author response:
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.

Reviewer comment (Experiments):
How does MBRCSL compare to other state-of-the-art offline RL algorithm on benchmark tasks?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
Can you formally characterize conditions under which model rollouts provide sufficient coverage for optimality of RCSL?

Author response:
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.

Reviewer comment (Experiments):
Can you incorporate offline constraints into RCSL policy learning?

Author response:
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.

Reviewer comment (Experiments):
Can you analyze sample efficiency of MBRCSL with limited suboptimal data?

Author response:
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

Reviewer comment (Experiments):
How many trajectories do you need before performing BC? I would imagine the performance of the algorithms depends a lot on this?

Author response:
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). 

| Task      | 1000   | 2000   | 3000   | 4000   | 5000   | 6000   | 7000   | 8000   | 9000   |
| --------- | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ |
| PickPlace | 0.1650 | 0.3875 | 0.4775 | 0.4525 | 0.4775 | 0.4475 | 0.4950 | 0.5125 | 0.4825 |

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.

Reviewer comment (Theory):
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. Unfortunately the experiments aren't particularly convincing in this respect.

Author response:
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.

Reviewer comment (Experiments):
Is end-to-end training of model and RCSL policy possible?

Author response:
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.

Reviewer comment (Experiments):
What are the most important architectural choices for dynamics model and RCSL policy?

Author response:
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.

Reviewer comment (Theory):
Can MBRCSL be extended to stochastic environments?

Author response:
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.

Reviewer comment (Theory):
The 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.

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
> **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:

| Tasks | Pointmaze | PickPlace | ClosedDrawer | BlockedDrawer |
| ----- | --------- | --------- | ------------ | ------------- |
| ratio | 14.3%     | 19.2%     | 25.7%        | 63.2%         |

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.

Reviewer comment (Theory):
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?

Author response:
> **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.

Reviewer comment (Theory):
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?

Author response:
> **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:

| Tasks | Pointmaze | PickPlace | ClosedDrawer | BlockedDrawer |
| ----- | --------- | --------- | ------------ | ------------- |
| ratio | 14.3%     | 19.2%     | 25.7%        | 63.2%         |

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.

Reviewer comment (Experiments):
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?

Author response:
> **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. 

| Task      | MBRCSL (ours) | MBDT     | MBBC     |
| --------- | ------------- | -------- | -------- |
| Pointmaze | **91.5±7.1**  | 79.8±4.6 | 83.0±7.5 |

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.

Reviewer comment (Writing):
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.

Author response:
> **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.