Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View.

Dear reviewer, We thank the reviewer for their detailed feedback. It seems like the reviewer's main concerns are the importance of theoretical contributions, introduction of a new hyperparameter and some missing experiments. We have revised Sections 5, 6 and:

1. Added Figure 8, in appendix B, to show that data augmentation works for a range of K values {10 25 40 50 75 100}.

2. Added Figure 11, in Appendix B which shows preliminary results for scalable SL algorithms built on top of our work can also achieve stitching generalization.

3. Added Appendix D.1, which shows that uses our definition of stitching, to prove that the optimal OCBC performs policy improvement over the behavior policy.

Do these clarification and new experiments and proofs address all the reviewer's concerns?

They indeed provided a formal definition of stitching generalization in Definition 1. However, it is not actually used anywhere.

We believe a formal definition of stitching generalization is useful because we use the definition to generate the datasets used to evaluate stitching. We revise section 5, para “Popular offline datasets do not evaluate stitching.” (see green text and equation), to clarify this. In Appendix D.1, we add a section which shows that the Bayes optimal OCBC policy that is learnt on a “stitched dataset”, where states-goals are sampled using the test distribution, is an improvement over the BC policy. In offline RL, policy improvement guarantees over the behavior policy are important because optimality guarantees are not generally realisable [1,2].

Weakness on Algorithm 1 (OCBC + Temporal Data Augmentation) regarding additional hyperparameter

In appendix B, Figure 8, we add a new experiment which ablates the value of the number of centroids on two environments (ant maze and point maze medium). We run experiments with K ranging between {10 25 40 50 75 100}. We can see that all values significantly improve the stitching generalization over the baseline.

Isn’t it a problem of trajectory-based sampling rather than transition-based sampling

OCBC algorithms also learn from individual transitions. For example in Figure 1, they will be trained on transitions (2,up,3) and (3,right,4). But they will never sample (2,up,4), which is out of distribution (has 0 support in the training distribution). There is no guarantee that a SL method will perform on points which are strictly out of distribution. But our paper’s main goal is to take initial steps to endow purely SL methods with these desirable properties of stitching-based algorithms like TD, while still retaining their simplicity and scalability.

It seems necessary to conduct additional Experiments on DT + temporal augmentation.

The reason we do not add data augmentation with DT is that because DT takes in a context of previous states and actions, the data augmentation method should cluster a history of states and actions, which is difficult to do computationally. We have added this reason in our updated paper (see Caption of Figure 4, text in green).

I also look forward to scalable version of their algorithm.

We believe that purely SL algorithms which learn state representations containing information about future outcomes, can possess desirable properties associated with TD learning while still retaining their simplicity and scalability. In Figure 11, in Appendix B, we include an experiment that shows using such representations as inputs for the OCBC policy, improves stitching generalization without the need of a distance metric.

There seems to be some minor typos.

Thank you for pointing these out, we have fixed these typos.

Dear reviewer, We thank the reviewer for their detailed feedback. It seems like the reviewer's main concerns are about the correctness of a theoretical claim (stitching vs generalization) and the limited experiments. We believe the theoretical results are correct (see below), and have revised Section 4, 5, and added a new proof in Appendix D.1 to clarify the results; we have also added new image-based experiments to bolster the empirical contributions of the paper. Below, we will describe these experiments in more detail and answer the specific reviewer questions Do these answers, revisions and experiments address all the reviewer's concerns?

Why is stitching a problem of generalizing to a different test distribution, and not a problem that occurs due to limited data from the training distribution?

Consider two datasets collected via two different sources:

Source 1 : Different data collecting policies $\beta_h$, with probability $p(h)$.

Source 2 : The BC policy $\beta$.

Both these datasets induce the same discounted distribution over states (Lemma 3.1), but the induced distribution over trajectories (state - goal pairs) is different in general, even after collecting infinite data from both sources (Lemma 4.1). This can be observed from Figure 1, where there are 4 possible trajectories. The BC policy $\beta$ induces a distribution that samples all 4 trajectories (Figure 1.c),. But the data collecting policies $\beta_h$ induce a distribution over only 2 trajectories (Figure 1.b). This holds true even if infinite data is collected from both sources. We agree that if we had access to the other 2 trajectories, then OCBC algorithms would be able to perform stitching generalization. Our claim is that these trajectories are out of training distribution, because we assume that the data is collected using Source 1. If we have access to all possible out-of-distribution trajectories, then of course generalization wouldn’t be an issue. Moreover, the perspective that, “if only you had data from the test distribution the problem would be solved, and hence it is a limited data problem” is incorrect when the train and test distribution are different. Because you will never have data from the test distribution by definition. And since the test distribution for stitching is different from the train distribution (Lemma 4.1), stitching is a generalization problem.

I think there are a few details that seem inconsistent throughout the paper, such as the fact that data is not the fundamental limitation vs popular datasets not evaluating stitching because they include most (s,g) pairs.

We want to clarify that this is not an inconsistency. In section 5, in para “Popular offline datasets do not evaluate stitching” we have added one sentence (in green color) to clarify this. The main claim of our paper is that: In offline RL settings, when you have multiple data-collecting policies, even if you collect infinite data from these policies, generally there will be state-goal pairs that would never be sampled together. At test time, we need to evaluate the algorithm on exactly these state-goal pairs to evaluate stitching. The problem with the original D4RL datasets is that they evaluate on the same state-goal pairs, which are visited by the same data collection policy. See Figure 10 (D4RL) vs Figure 3 (ours) to see this difference.

experimental evaluation seems limited overall, and somewhat preliminary.

We have added an image-based version of the point maze environment which uses a top-down image of the maze instead of the low-level x,y coordinates (see Figure 9). Figure 9 reaffirms that OCBC algorithms are unable to stitch. And despite the high dimensionality, our data augmentation does perform stitching. We would like to add that our experiments on Ant locomotion tasks, which are very popular in offline RL literature [1,2,3], and are considered to be challenging (much harder than the D4RL’s locomotion tasks).

Stitching as goal-state pairs : My understanding of the stitching property is that it is more general

We agree that the term stitching is used more generally. At the start of section 4, we have added a paragraph (green text) that talks about different properties which are associated with stitching. Our definition indeed covers all the existing trajectories that can be re-combined using dynamic programming. In the offline RL setting, dynamic programming methods compare all actions (be it from different trajectories) that were taken in a particular state. By definition, these actions come from the distribution of the BC policy. This is exactly the distribution used to test for stitching. Our definition can be trivially extended to the reward-case, by substituting state-returns pairs (similar to the Decision Transformer paper) instead of the state-goal pairs.

We thank the reviewer for their comments and further discussion on the work.

We believe the "relevance of this problem setting" is very important to the RL, supervised learning, and planning communities.

A few questions remain around the relevance of this problem setting, because the stitching property is not one that has been investigated before given the state of the very common offline RL benchmark (D4RL).

The D4RL paper [1] has stated multiple times that testing for stitching is important (stitching is mentioned 6 times, See section 4,5 of [1] ). This raises an important question of when and whether stitching occurs. The first important contribution of our paper is to give a concrete definition of stitching as a form of generalization. Given this relation, we build on D4RL datasets to explicitly test for stitching. Indeed, our experiments reveal that OCBC algorithms which perform well on the original D4RL datasets [8], cannot perform stitching. We open source our code and the datasets needed to test for stitching. Arguably, stitching is the key property that helps offline RL algorithms find better solutions than the behavior policy. See appendix D.1, where we prove that the stitching is required to show one-step policy improvement over the behavior policy.

Given the recent advent of using large supervised learning transformer models for many problems in RL, it is important to ground these approaches. This has become a very popular and relevant topic in RL [2,3,4,5,6,7,8,9] as well as in SL [10,11,12,13]. Our work shows both theoretically and empirically, why OCBC algorithms, even with larger transformers and larger datasets will not perform stitching. Our concrete definition of stitching, and open source datasets will give future researchers a chance to improve OCBC algorithms. While our work shows that current OCBC algorithms do not stitch, we also propose a simple data augmentation which outperforms the baselines on all tasks on state-based (Fig 4) as well as image-based tasks (Figure 9,10). In the light of these contributions, we believe that many RL researchers at ICLR 2024 would find our work useful and interesting. We kindly ask the reviewer to reconsider the paper in light of the revisions and clarifications.

[1] D4RL: Datasets for Deep Data-Driven Reinforcement Learning

[2] Imitating Past Successes can be Very Suboptimal

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

[4] You can’t count on luck: Why decision transformers and rvs fail in stochastic environments

[5] Dichotomy of Control: Separating What You Can Control from What You Cannot

[6] Upside-Down Reinforcement Learning Can Diverge in Stochastic Environments With Episodic Reset

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

[8] RvS: What is Essential for Offline RL via Supervised Learning?

[9] Is Conditional Generative Modeling all you need for Decision-Making?

[10] Faith and Fate: Limits of Transformers on Compositionality

[11] Location attention for extrapolation to longer sequences

[12] Unveiling transformers with LEGO: a synthetic reasoning task

[13] NaturalProver: Grounded mathematical proof generation with language models.

Both the changes to the paper and the answers to my questions clarify my misunderstandings. My new understanding is that stitching generalization is different from iid generalization because there is a latent structure in the (state,goal) pairs seen during training and those that are relevant for testing. Thus, the distribution is not identical nor independently distributed. The edge-case that I described, where all states and goals are present in the dataset, would not require stitiching generalization (and presumably, the supervised learning approaches would work here).

I appreciate your reply. A few questions remain around the relevance of this problem setting, because the stitching property is not one that has been investigated before given the state of the very common offline RL benchmark (D4RL). These is value, however, in pointing out this property and I think this work need not answer that question. I will update my score (5 -> 6) to reflect this exchange.

Dear reviewer, We thank the reviewer for their detailed feedback. It seems like the reviewer's main concerns are regarding a clear definition of stitching, a theoretical justification for the data-augmentation, and the difficulty of obtaining a good distance metric. We have incorporated these concerns using new updates to the papers and new image based experiments:

1. We add a new paragraph at the start of Section 4 (green text) to explain how stitching is colloquially associated with different properties of TD learning. We add that we focus on one of these properties as a clear difference between TD and OCBC methods and call it stitching generalization, which we concretely define in Definition 1.

2. In Figures 9 and 10, in Appendix B, we perform additional experiments with an image-based version of the point mazes. We show that performing data augmentation works with images as well, where a good distance metric is not available.

3. In section 5, we added a paragraph “Theoretical intuition on temporal data augmentation”, to provide some theoretical justification for data augmentation.

4. In Figure 11, in Appendix B, we include a preliminary experiment that removes the use of a distance metric, by learning representations. We believe that this approach is a promising direction for the future work that builds on our paper.

Below, we will describe these updates and experiments in more detail and answer the specific reviewer questions Do these answers, revisions and experiments address all the reviewer's concerns?

In the RL context, similar states do not mean they have similar outcomes

We agree with the reviewer that having a good distance metric is challenging and have already acknowledged the limitation in the main paper. But empirically, we show that this limitation might not always matter. We add two sets of experiments ( appendix B.2, Figures 9 and 10 ) on image-based pointmaze tasks and show that K-means using L2 distances between images can perform well, significantly improving over the OCBC baselines on all tasks. We would like to add that for all of our maze experiments, if the cut-off distances for K-means clustering was too large, then it would group together states on different sides of the walls like you suggest. But our empirical results (Figures 4,9, and 10) show that using K-means works well.

"Does more data remove the need for augmentation" is not very surprising given the construction of the dataset

In this work, we aim to understand and evaluate the stitching properties of OCBC algorithms. The experiments are designed to evaluate this question and support the claim that data augmentation isn't just about inflating the size of the dataset. This claim (and verification) is important because in today's world of large models, many researchers might otherwise be tempted to throw large amounts of data at this problem, without avail.

Q1: Any suggestions for a reasonable distance metric for the states?

We believe that learning state representations such that the distance between these representations contains information about the similarity of their outcomes, is a scalable and promising direction of future work [1,2]. This will solve the exact problem that you correctly pointed out. In Figure 11, in Appendix B, we include an experiment that shows using such representations as inputs for the OCBC policy, improves stitching generalization without the need of a distance metric.

Q2: Are there any theoretical justification for the temporal augmentation?

While exact theoretical guarantees for temporal data augmentation will depend on how good the distance metric is, we can think of the temporal data augmentation as trying to diversify the dataset by sampling from $\pi^{\beta}_+(g \mid s)$. Since this is the test distribution for stitching, data augmentation should improve stitching generalization. In section 5, we add a paragraph, “Theoretical intuition on temporal data augmentation” to intuitively explain how data augmentation leads to a better policy.

Minor Thank you for spotting these errors, we have fixed them in our updated paper.

[1] DATA-EFFICIENT REINFORCEMENT LEARNING WITH SELF-PREDICTIVE REPRESENTATIONS [2] Contrastive Learning As a Reinforcement Learning Algorithm

the additional explanation (first paragraph of Section 4)

This is meant to be an intuitive introduction on how stitching is colloquially associated with various properties in the literature. In Definition 1, we give the precise definition of stitching generalization. If you think there is something about the definition that isn't rigorous, please let us know. We'd be happy to look into it.

the algorithm development lacks theoretical justifications

The data augmentation we propose does have theoretical backing. In Figure 2, in the limit of finest clusters, i.e, as the size of the blue circle tends to zero, temporal data augmentation is equivalent to sampling from the Markov chain s -> w -> g. This shows that in the limit of the finest clustering, temporal data augmentation will sample cross trajectory state-goal pairs. In practice, we show that K-means works well across various values of K (See Figure 8).

images can still be favorable to L2 distance because similar images

We agree that there can be image based tasks where it will be difficult to make data augmentation work. But the image based maze tasks are challenging, and adding the temporal data-augmentation (5 lines of code) improves the performance of baselines by more than 2 times on all the 3 image based tasks.

We kindly ask the reviewer to reconsider the paper in light of the revisions and clarifications.

Dear reviewer, We thank the reviewer for their detailed feedback. It seems like the reviewer's main concerns regarding the assumption of a distance metric, especially in high dimensional states, and the inclusion of contrastive representation-based baselines. We address these concerns by adding two new experiments and updates:

1. In Figures 9 and 10, in Appendix B, we perform an additional experiment with an image-based version of the point mazes. We show that performing data augmentation using k means on images performs much better than the baseline. Although we agree that finding a good distance metric is difficult, which we have already mentioned in the paper, our new results using images and the existing ant maze tasks with 29 dimensional state spaces, suggests that K-means with simple L2 distance can empirically still perform great.

2. In Appendix D.2, we prove that contrastive RL will not learn Q values for states, goals that are used to test for stitching.

3. In Figure 11, in Appendix B, we include a novel version of RvS which takes in input the contrastive representations of states [1]. Just like the reviewer suggested, we see that these representations are indeed able to improve the stitching abilities of OCBC methods. We believe that this approach is a promising direction for the future work that builds on our paper.

Below, we will describe these updates and experiments in more detail and answer the specific reviewer questions Do these answers, revisions and experiments address all the reviewer's concerns?

The proposed stitching generalization describes how "far away" different goals

The information of how far away goals are from a particular state is not directly available in the distance metric. We do not make this assumption, and it is not even required. We only need a distance metric to make sense on very small length scales, that is it should identify states which are very close together. This is easier than describing how “far away” goals are from a given state (Q value of reaching a goal). To actually learn to actually select actions that help you reach unseen goals, the temporal data augmentation we propose is important.

The experiments are limited to Antmaze

To incorporate your suggestion, we include two sets of experiments with the image based version of point maze tasks (umaze, medium and large). In figure 9 and 10, in appendix B, we see that applying clustering on images and using it for our data augmentation outperforms the baselines on all tasks. It also matches the performance of the state based augmentation in the large maze (Figure 10).

Missing baselines: Contrastive-based learning is also studied for goal-reaching problems as one kind of supervised learning problem,

In Appendix D.2, we add a proof (based on Lemma 4.1) which shows that contrastive RL [1] will not estimate the Q function for state-goal pairs used to test for stitching. We will include the actual empirical results in the camera ready version, as we need more time to adapt their code (which is written in a different framework) to our settings.

it might also learn some representation more useful than the baseline included in the paper.

According to your suggestions, we performed an ablation experiment which uses contrastive representations [1] as inputs to the OCBC policy on point maze tasks (umaze, medium and large). In Figure 11, in Appendix B, we see that these representations are indeed able to improve the stitching abilities of OCBC algorithms. We have added Appendix B.3, which notes that learning contrastive representations for OCBC algorithms is a promising direction for future work.

The optimal augmentation would be some metric that describes "how "far away" different states are in the sense of transitions" almost perfectly.

This proposed metric could work well without our augmentation method, however it is more than what is needed to improve performance. We do need a distance metric to cluster close-by states. Intuitively, we want a distance metric that can classify states as very close or not. See for example, figure 2 where the distance metric does not tell us how far $\tilde{g}$ is from $s$. It only tells us that $w$ and $g$ are close-by states.

[1] Contrastive Learning As a Reinforcement Learning Algorithm