When does return-conditioned supervised learning work for offline reinforcement learning?

Abstract: Several recent works have proposed a class of algorithms for the offline reinforcement learning (RL) problem that we will refer to as return-conditioned supervised learning (RCSL). RCSL algorithms learn the distribution of actions conditioned on both the state and the return of the trajectory. Then they define a policy by conditioning on achieving high return. In this paper, we provide a rigorous study of the capabilities and limitations of RCSL something which is crucially missing in previous work. We find that RCSL returns the optimal policy under a set of assumptions that are stronger than those needed for the more traditional dynamic programming-based algorithms. We provide specific examples of MDPs and datasets that illustrate the necessity of these assumptions and the limits of RCSL. Finally, we present empirical evidence that these limitations will also cause issues in practice by providing illustrative experiments in simple point-mass environments and on datasets from the D4RL benchmark.

Reviews and author responses:

 Thanks for engaging in the discussion!

1. It is still not clear to us why the reviewer thinks that RCSL is just BC. This does not seem to be substantiated and there are key differences, namely the ability to condition on returns and outperform the behavior policy.

2. We want to reiterate that there is no need for the function class to be fully expressive. We agree that using tabular function approximation with full realizability of all policies in an adversarially constructed problem RCSL requires a larger function class than the policy in a DP algorithm. But no one is advocating for this function class! These are methods intended to be used with function approximation and rewards are not naturally represented as discrete variables. In practice, we use neural nets of very similar expressivity for all methods. RCSL depends on approximation error as defined in assumption 3 of corollary 3 and this only requires approximation of the optimal policy not the full policy class (for example epsilon could be zero even if the function class only has one function if it is the optimal one). The reviewer instead is discussing the assumptions 1 and 2 in section 5.1 that are necessary for guarantees about the DP algorithm. In contrast, those DP assumptions require realizability and completeness in the space of Q functions rather than policies. The assumptions are not directly comparably, but we would argue that the DP assumptions are stronger since they require both a realizability assumption and a completeness assumption while RCSL only requires a notion of realizability. Of course it is possible that the RCSL map from S x R -> A is more complicated that the DP map from S x A -> R, but we don’t think there is a clear argument either way in general and think the two function classes should be viewed as having similar complexity in theory and in practice.

Please let us know if anything is still unclear or if you would be willing to raise your score.


 Thanks for engaging in the discussion!
There seem to be two fundamental issues raised by the reviewer: (1) the title and (2) the experimental results.

1. On the title: there seems to be some confusion caused by our title. We ask the question: when does RCSL work for offline RL? We answer this question by our main contribution which is the theory that gives precise conditions under which RCSL works and examples where it fails. However, the conditions are in fact quite strong (and necessary) so that in practice we often expect RCSL to fail, which we see in the experiments. We want to make clear that our main claim is not that RCSL works, but that it can work under specific settings while failing under others. Would the reviewer be more satisfied if we changed the title to something like “Analyzing RCSL for offline RL”?

2. We have added some additional experimental results and discussion to a new Appendix A for several different types of datasets in the three D4RL environments from the main text for a total of 9 experiments across 6 algorithms in D4RL now. This way we have coverage of multiple dataset types in each of three substantially different environments from the D4RL benchmarks (one from each of the classes of environments). The results confirm those that were already in the paper, mainly highlighting that RCSL is usually beaten by DP-based algorithms (as we might expect in cases with poor return coverage as the theory suggests), but there are situations where this may not be the case (specifically cases with high-quality trajectories in the dataset, but poor coverage over the state space, as we may expect from the theory). Hopefully these results are sufficient to convince the reviewer to raise their score. We again want to emphasize that our primary results are theoretical and these experiments are not our main contribution, they simply provide one demonstration of how to interpret the insights from the theory.

Let us know if this response was satisfactory to raise your score or if there is anything else that is still unclear. 


 I think all my concerns except for the evaluations on stochastic examples were well addressed in the rebuttal. I think this paper is worth being accepted by NeurIPS, but I will keep my score unchanged considering the current experimental results.



 Thanks for the detailed responses from the authors.

The reviewer appreciates the honesty of the authors. As discussed in the main review, the theoretical bound for DP-based methods and that of RCSL are not directly comparable. Also, it is better to put the theoretical results of top-% BC into the main text.

Whereas the reviewer still deems that it is vital to figure out why RCSL fails while DP-based method works from a theoretical view, as this paper discusses "when does RCSL work for offline RL". The authors are also suggested to fulfill it through empirical experiments. While unfortunately, as the reviewer commented in the main review, the experiments are not sufficient. The authors can take some more datasets from MuJoCo or Adroit and consider the datasets from kitchen.

When commenting "RCSL fails/underperforms baselines on 3 out of 4 datasets", the reviewer does not mean that the authors need to make RCSL beat DP-based methods. The reviewer wants to challenge the authors on the key claims of the paper, i.e., the title. With only one tested task where RCSL beats DP-based method, can we really say that the evidence is enough to convince readers on the proposed claims in the paper?

 Dear authors,

Thank you for your detailed responses. 

(1) I understand that "this algorithm" has never been studied before, but my point is that in this restrictive setting, "this algorithm" is just BC, so it's not a super novel result in its own right that RCSL, given enough data, converges to optimality. 

(2): The dimension blow-up arises from history, not an additional state observation. Consider your reward to be discrete over R integer values, and the environment is H step. Then, to have a look-up table storing all possible policies, you would need to increase the size of the look up table from |S||A| to roughly the order of |S||A||R||H|. 

Now, in the non-tabular setting this paper focuses on, the size of the function class is absorbed into the epsilons in Assumption 1 and 2. However, to obtain the same epsilon, the function class for RCSL policies necessarily is bigger than that of regular policies. This key detail is not compared when comparing the regret bound to DP algorithms.

While the authors state in the rebuttal that for RCSL algorithms, the function class does not need to be able to represent all policies, Assumption 1 and 2 in Section 5.1 appear to contradict this point (for the bound to be reasonably meaningful).

Finally, I understand that in practice, the implementation would not be so different. But the authors have positioned this paper to be a theory paper, so this point should be properly addressed in theory. 

 First, we would like to thank the reviewer for their thoughtful review and positive evaluation of our paper!

The reviewer essentially raises two main areas for improvement: (1) more experiments and (2) a more explicit characterization of lower bounds.

1. We agree with the reviewer that more experiments could always make our point stronger. We are working on the following experiments based on the reviewers advice to add for the camera ready: (1) we will run on more of the D4RL datasets with different behavior policies, (2) we will add a variety of behavior policies rather than just reward functions in the toy example to see how performance degrades/improves as the behavior changes, and (3) we are also planning to add a stochastic toy environment in response to reviewer a7eE that may also be of interest.

2. We agree that formalizing the lower bounds could sharpen the contribution. This should not be too difficult as we already have derived the hard instances and presented them in the figures, but we will add a more formal treatment of lower bounds for the camera ready when we are allowed a longer page limit. The reviewer also suggests that we add more discussion of limitations, which we will add for the camera ready with the longer page limit. 

If anything still seems to be lacking or unclear, please let us know.


 First, we would like to thank the reviewer for their thoughtful review and positive comments about the originality, significance, and quality of our paper. If we are able to resolve some misunderstandings, we hope that the reviewer will be willing to raise their score.

Weaknesses:

a) Indeed the reviewer is right that some of the figures are slightly blurred, we will fix this in the revision. Thank you for catching this!

b) We are happy to make the requested change to line 199 in the revision, but it is not clear to us why the first version is not acceptable to the reviewer. This seems to us to be a fairly standard way to write papers in machine learning, so we are not sure where else to make changes to satisfy the reviewer. If you are able to raise more specific instances  that are problematic, please let us know and we can make the changes.

c) We agree with the reviewer that the known bounds for the two algorithms are not directly comparable since they make different assumptions. In this section we hope to provide some contrast that weighs the relative strengths of the algorithms and we think that our analysis provides solid high-level guidance comparing the assumptions needed for the different algorithms. However, we agree with the reviewer that these bounds are not perfect and will not tell us precisely when one algorithm will work and the other will not in practice because they are not tight in every instance and they all rely on assumptions that may not be satisfied. However, we think the bounds and the discussion are instructive to give general guidance about the types of factors worth considering (e.g. stochasticity of the environment, coverage of optimal behavior in the dataset, coverage of optimal state distributions in BC, or expressivity concerns in DP).

d) We want to make it very clear that we are not arguing that RCSL is better than the alternative algorithms we implement in the experiments. The purpose of our paper is to examine and analyze how RCSL compares to more well-established methods like BC and DP. To that end the toy experiments and examples provide the clearest evidence of the claims of the paper and then we consider the benchmark datasets to see how the insights can be applied to problems where the exact correspondence to the theory is less clear and where we would expect to see less clear cut results. We chose a sample of datasets from D4RL that have highly different properties to see what happens in different cases since we think there is not much insight to be gained from the minor variation across most of the D4RL datasets. However, we can run the experiments on more datasets for the camera ready, just let us know which datasets would be useful to see results from. 

Questions:

a) We make a brief comment about the connection between RCSL and distributional RL to aid readers in understanding the connection between the two. Both algorithms deal with the same object which is the distribution over future returns. In distributional RL this distribution over returns is modeled explicitly, while RCSL implicitly uses this same distribution to define its policy. We can remove the comment if the reviewer finds it too confusing, but we thought it provided some useful context.

b) Indeed, top-% BC suffers from the same issues as RCSL in stochastic environments since it filters using the stochastic returns. Since this algorithm is not the focus of the paper we did not explicitly create these example MDPs for top-% BC, but we can add them in the revision if the reviewer thinks it would be useful. 

c) RCSL training is simply supervised learning and as such is as easy to train and debug as supervised learning is. We agree that when using complex models like transformers this is not trivial, but supervised learning is widely regarded as easier to train than DP-based methods, this is all we are trying to claim. 

d) RCSL has the potential to be beneficial in cases where the dataset already contains some high-return trajectories, but where the dataset does not have good coverage over the state space (as would be necessary for DP-based methods). We see this somewhat in the results on the human-pen-v1 dataset. However, we agree with the reviewer that our results are largely negative about the RCSL algorithm. Again, we want to reiterate that we are not trying to claim that RCSL is the best algorithm. Rather, we are analyzing a recently created but not well understood approach. As to the questions of stochasticity, we have shown that RCSL as presented has fundamental issues. Trajectory transformer is a different method that is essentially model-based RL and as such can deal with issues like stochasticity. It is beyond the scope of the paper to comment more deeply about how to resolve these issues. 

Hopefully these responses can clear up some misunderstandings and the reviewer is willing to increase their score. If anything is still unclear that may impact the score, please let us know! 

 First, we'd like to thank the reviewer for their thoughtful review and largely positive assessment of the paper. We will respond to the weaknesses first and then the question.

Weaknesses:

1. It is not exactly clear to us what the reviewer means here by an augmented state space. In some sense the idea of RCSL is to add the future return to the state during training, but this algorithm has never before been analyzed to our knowledge. If the reviewer has a specific result in mind, we would be happy to compare it to our work more directly. We also want to clarify that while the dynamics are (nearly) deterministic, the behavior policy that generates the data can be highly stochastic, so the return itself is not at all a deterministic quantity, even when conditioned on the state.

2. We don't see any obvious issue with the comparison of function class sizes between RCSL and DP. There seem to be two points of misunderstanding here. (1) It seems that the reviewer is taking issue with our presentation of RCSL as history-conditioned rather than just conditioning on the most recent observation (and this is where the exponential dependence is arising). We presented RCSL as history conditioned to be fully general and to capture the Decision Transformer work, but in Markovian environments, all of our analysis holds for simple state-conditioned policies without any history. Moreover, in non-Markovian settings even standard DP algorithms will have to condition on the entire history. (2) It seems that the reviewer is assuming that our policy class must be fully representative of all possible policies (as in a full lookup table of all states and actions). This need not be the case, and for our positive results we only need the policy class to be able to realize the optimal policy to compete with it. This is a strength of RCSL compared to DP, since DP algorithms require stronger assumptions on the Q function class beyond realizability. Finally, we just want to note that in practice there is essentially no difference in the actual neural networks we use to represent the RCSL policies versus the DP policies and Q functions. For example the RvS policy, TD3+BC policy and Q function, and IQL policy and Q function in our implementation have the same base architecture and are not history dependent (with only minor differences in input space: (s, g) for RvS policies, (s) for DP policies, and (s,a) for DP Q functions). The decision transformer does indeed have a different architecture, but this is somewhat orthogonal to the point we are making about the underlying loss functions that guide our analysis.

Questions:

1. We agree with the reviewer that section 6.2 is not a particularly novel experiment since we are just applying prior methods to a standard benchmark. We simply include it for completeness of our discussion and do not try to claim that it is our core contribution.

We hope that these responses clarify our contribution and encourage the reviewer to raise their score. If there are any remaining questions about our response that may impact the score, please let us know!



 First, we'd like to thank the reviewer for their thoughtful review and largely positive evaluation of the paper. Below we respond to each of the questions the reviewer raises and then the weaknesses (although there is some overlap).

1. Perhaps we were not clear enough in our explanation of how RCSL works and we will try to clear this up here. There is no mismatch between g(tau) and f(s), they are both meant to represent the return over trajectories. At test time, we don't have access to the future, so RCSL conditions on a desired value for the future return of the trajectory from the current timestep onwards and we denote this desired return as f(s). We also want to emphasize that we are not proposing the RCSL algorithm as our contribution, rather we are providing an analysis on prior work that introduces the algorithm and pointing out the flaws and strengths of the approach.

2. Indeed N is dataset size, thanks for catching our omission! We will add the definition to the revision.

3. Initially we omitted stochasticity from the experiments since we already had simple examples where stochasticity provably causes RCSL to fail. However, the reviewer raises a reasonable point that a more full characterization of the issue in a toy experiment could improve our argument. We are working on implementing this now, and hope to have results soon.

The only weakness that the reviewer lists is to add an experiment to address question 3, which we are working on now. Hopefully this response clarifies things enough for the reviewer to increase their score. If there are any other weaknesses that need to be addressed to improve the score, please let us know.


 This paper studies the capacities and limitations of return-conditioned supervised learning (RCSL) methods and illustrates their relationship to behavior cloning (BC) and dynamic programming (DP). The main conclusion is that RCSL methods are theoretically incapable of achieving high performance when the underlying environment is non-deterministic and the return coverage is insufficient. These findings are both intuitive and theoretically sound. The novelty of this paper lies in the theoretical part. The authors start with the setting with infinite data and a fully expressive policy and then extend the derived theorems to a more practical setting with finite data and a limited policy class. The theoretical derivations are solid and the theorems are easy to understand. The core conclusion is that RCSL is fundamentally limited when the environment is stochastic and the return coverage is insufficient, which is intuitive and theoretically sound. However, from my perspective, I think the authors should strengthen the experimental parts. For example, to verify that RCSL is indeed limited to near-deterministic environments, the authors should conduct experiments in environments with different stochasticity. 1. In the training phase, RCSL minimizes the empirical negative log-likelihood loss described by Equation (1), while the test-time policy is defined as Equation (2). It seems that there is a mismatch between the training-time policy and the test-time one: the training-time policy is conditioned on the return of an entire trajectory $g(\tau)$ but the test-time policy is conditioned on a condition function $f(s)$ which varies with the states in a trajectory. The authors should explain it. Besides, will this discrepancy affect the conclusion?
2. The authors did not pre-define the variable $N$, which firstly appears in line 182 and Equation 9. (I guess it refers to the dataset size)
3. Though the authors claim that RCSL is incapable of learning high-performance policies when the environment is stochastic, which is described by Theorem 1 and explained by Figure 1, I think the authors may construct toy MDPs to provide experimental evidence in the experimental part.
 It would be better if the authors could conduct experiments to support their theoretical claims. For example, experiments investigating the return coverage and stochasticity of the dataset and the environment, respectively.

 This paper theoretically analyzes a new popular class of RL algorithms, referred to as Return-Conditioned Supervised Learning (RCSL). The paper theoretically shows that this class of method requires stronger assumptions than standard DP-based approaches for learning the optimal policy. The paper shows that RCSL method requires near-deterministic dynamics as well as a priori knowledge of the optimal conditioning function to perform well, and construct examples where the stated assumptions are necessary for empirical performance.  Strength:
1. This paper is solid in all four axis: originality, quality, clarity, and significance. It is the first paper to provide a theoretical analysis of a new class of practically relevant RL algorithms. The exposition is clear, and the theoretical results are mostly proper (though, I did not check the supplement fully). The experiments are well thought out, and complement the theoretical analysis. Given that this class of algorithm has received significant attention from the RL community, this work is timely and relevant. 

Weakness:
1. It's unclear that under deterministic dynamics, RCSL can recover the optimal policy, is a novel result. Under this assumption, the return can be practically treated as a new state-dimension, and RCSL reduces to Behavior Cloning with an augmented state space.
2. The dependence on the size of the policy class needs to be clarified when comparing RCSL and DP-based algorithms in Section 5. Even in the tabular case, just to represent the policy, it would require a lookup table that is exponentially larger than what is required to represent a standard policy for DP approaches (i.e, |S||A|). For this reason, it's not clear that the comparison is super meaningful. 
 1. It's unclear what insight Section 6.2 adds; I believe similar findings have already been reported in prior papers, for example, https://arxiv.org/abs/2112.10751. 

I am willing to increase my score if the authors can address my qualm about the theoretical analysis in the above section. Overall, this is a sound paper.  Yes. 

 This paper discuss the conditions that return-conditioned supervised learning (RCSL) work for offline reinforcement learning (RL). This paper mainly contributes theoretically on the optimality of the RCSL algorithms and the sample complexity of the RCSL algorithms. Under some strong assumptions, they present theoretical guarantees for RCSL given both infinite data and finite data. Furthermore, empirical evidence on failure and success of RCSL methods are provided. **Originality and Significance**: this paper is quite original, which theoretically studies a very new topic in RL. The conclusions in this paper is of importance to the RL community, and can benefit the progress of offline RL

**Quality**: this paper has a good quality overall. The theoretical results seem sound to some extent. The reviewer randomly checked some proofs in the appendix (but not all of them), and they seem correct, despite some minor errors in some of them. For example, in appendix Eq. (56), the reviewer thinks it ought to be $\epsilon (\frac{1}{\alpha_p}+1)$ instead of $2\epsilon (\frac{1}{\alpha_p}+1)$. The theoretical contribution of this paper is clear. While the empirical experiments are kind of disappointing

**Clarity**: this paper is of medium clarity. The reviewer likes the way the authors tell a story like this. While the figures are of poor quality, which negatively affect the clarity of this paper. Some figures, e.g. figure 3 and 4, are sort of blur.

**Summary of Strengths**:

a) the discussed topic is important and interesting. It is vital to figure out when and under what conditions will RCSL work for offline RL

b) the theoretical contribution of this paper is of importance to better understand "when return-conditioned supervised learning work for offline reinforcement learning"

c) the toy examples presented in the paper is quite helpful to understand why RCSL fail in stochastic environments

**Weaknesses**:

a) some of the figures are blur.

b) there are some grammatical issues, e.g., line 199 "To see why this is the case, consider the MDP illustrated" --> "To see why this is the case, we consider the MDP illustrated". Please do not use colloquial expressions in submitted papers. While this is just a random example. Please double check the paper to ensure that there are no grammatical issues.

c) the discussion lacks depth. For example, in Section 5.1, the authors compare the sample complexity bound of RCSL and dynamic programming (DP). The reviewer cannot tell what is the point here, since the basic assumptions for these bounds are different. Whereas the bound for top-%BC are omitted in the main text. Meanwhile, the discussion in Section 5.1 does not reveal why RCSL fails while DP-based methods work.

d) the empirical evidence is expected to be enriched. Currently, RCSL is only tested on some selected datasets, e.g., 1 D4RL MuJoCo dataset, 1 D4RL Adroit dataset, and 2 D4RL Antmaze datasets. RCSL fails/underperforms baselines on 3 out of 4 datasets. The reviewer agrees that MuJoCo/Antmaze datasets come from "stochastic" environments. While there is only one single test on datasets with narrow distribution (e.g., human). The reviewer cannot say that it can convince readers on the proposed claims in the paper. The reviewer has the following several questions:

a) why discussing the connection of RCSL to distributional RL? They actually lie in quite different categories.

b) the theoretical results in appendix (on top-%BC) also involve the assumption of "near determinism". Then will top-%BC fail on stochastic environments? If we include all the data, i.e., top-100% BC or vanilla BC, are the derived bounds still meaningful? Any discussion here?

c) the authors claim that DP-based methods are often "unstable and difficult to debug" (line 371-372), then are RCSL methods easy to debug and exhibit stability during training? Why RCSL methods are easy to debug? It is hard to say that decision transformer (DT) is easy to tune and debug.

d) what is the benefits of RCSL on earth? Since IQL, a DP-based method, shows very good performance over all of the tested datasets. Does there exist some methods that can make RCSL work in stochastic environments? The authors mention that trajectory transformer works on MuJoCo datasets as it involves planning, will planning breaks this dilemma? N/A

 The paper considers return-conditioned supervised learning (RCSL), a class of methods that learn a return-conditioned policy from offline data and recover the evaluation policy by conditioning on a given desired return. It then analyzes these method theoretically, bounding the regret of RCSL policies under the assumption of close-to-deterministic transitions&rewards and return-coverage of the behavior policy. Moreover, the paper provides examples of MDPs where RCLS is not able to recover the optimal policy, which demonstrate that the additional assumptions are required. Lastly, in the experimental evaluation the paper shows on a toy-example scenarios where dynamic programming has advantages over RCSL and also evaluate performance on the D4RL dataset. The paper is clearly written, easy to follow, and studies a class of methods that has been of interest lately. The theoretical analysis clearly illustrates when we can expect RCSL to work and the example-MDPs are useful to understand when it doesn't. The comparison to dynamic programming and behavior cloning is interesting, but comparing lower-bounds would probably be more convincing. The experiment section illustrates the main points of the paper, but is relatively limited in scope. Overall I think this is a well-written and interesting paper. It would perhaps be useful to explicitly state a lower-bound on performance given that the introduction is largely motivated from that perspective. It would be nice to expand the experiment section a bit further: For the toy-example and ablation study on how performance changes depending on the behavior policy would be useful and more D4RL baselines could be run. For the latter it would be particularly interesting how performance compares on different behavior datasets for the same environment. While the theory clearly states its assumptions, limitations of the analysis are not discussed.

Meta-review: This paper theoretically analyzes a new popular class of RL algorithms, referred to as Return-Conditioned Supervised Learning (RCSL). The paper theoretically shows that this class of method requires stronger assumptions than standard DP-based approaches for learning the optimal policy. The paper shows that RCSL method requires near-deterministic dynamics as well as a priori knowledge of the optimal conditioning function to perform well, and construct examples where the stated assumptions are necessary for empirical performance.

This paper is solid in all four axes: originality, quality, clarity, and significance.  All reviewers advocated for acceptance.



Decision: Accept