Paper: Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View. (ICLR-2024)

Reviewer comment (Theory):
**Theoretical contribution is a bit marginal.**  
The authors claim that they “provide a theoretical framework for studying stitching.” They indeed provided a formal definition of stitching generalization in Definition 1. However, it is not actually used anywhere. If it is not applicable to deduce any theoretical guarantee, what is the purpose of proposing such definition? At least for the counterexample proposed in Lemma 4.1 and depicted in Figure 1, the authors should be able to provide any kinds of implications on the (lower bound of) stitching generalization. This will definitely strengthen the theoretical contribution of the paper. If it is not possible, please kindly explain the reason/context.

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

Reviewer comment (Theory):
**Isn’t it a problem of trajectory-based sampling rather than transition-based sampling that it is impossible to implement stitching with SL-based methods?**  
At first glance, I thought that the statement “stitching generalization is not the same as i.i.d. generalization” is not quite the problem of supervised learning itself. Let’s take a look at the counterexample in Figure 1. If we break down the trajectories collected by policies $\beta_{h=1}$ and $\beta_{h=2}$ into transitions, can’t even a supervised method learn something from a virtual trajectory, say, 2-3-4, which was not exactly in the dataset, by only i.i.d. sampling the transitions? To rephrase my question, can’t a SL-based method learn stitched behaviors by transition-based sampling or so-called *bootstrapping*? I don’t think this would be exactly the same as the proposed temporal augmentation method. Please correct me if I’m wrong.

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

Reviewer comment (Experiments):
**Weakness on Algorithm 1 (OCBC + Temporal Data Augmentation) regarding additional hyperparameter**  
It uses clustering algorithm, so it has an additional hyperparameter, $k$ for k-means algorithm. To the best of my knowledge, the main paper does not discuss on the effect of the choice of such hyperparameter, which might be difficult to properly tune.

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

Reviewer comment (Experiments):
**It seems necessary to conduct additional Experiments on DT + temporal augmentation.**  
Although the paper empirically proves that it can enhance the data utilization of *RvS* algorithm with their proposed augmentation technique, it does not provide any results on ‘’DT + temporal augmentation’’ combination. This experiment will strengthen the empirical contribution of the paper. If it is not applicable, please kindly explain why.

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

Reviewer comment (Writing):
**There seems to be some minor typos.**  
- pg 1. Kaelbling (1993) $\rightarrow$ (Kaelbling, 1993)  
- pg 3., Equation (1). $\mathbb{E}\_{s\sim p_0(s_0)} [p_t^{\pi}(s_{t+} =g \mid s_0=s)]$ $\rightarrow$  $\mathbb{E}\_{s\sim p_0(s_0)}[p_+^{\pi}(s_{t+} =g \mid s_0=s)]$  
- pg 3. “… and then sampling a trajectory from the corresponding policy $\\{\beta(a\mid s, h)\\}$.” $\rightarrow$ $\beta(a\mid s, h)$

Author response:
> There seems to be some minor typos.

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

Reviewer comment (Theory):
In Definition 1, does the performance function $f$ have any relavance to the objective $J$ defined in Equation (3) or the maximum likelihood objective defined in Equation (6) and (7)?

Author response:
> does the performance function any relevance to the objective 

The performance measure can depend on the practitioner. It can be the time required for the agent to reach the goal (shortest path) or it can be the probability that the agent reaches and stays in the goal (goal reaching). So the objective of OCBC algorithms and the performance measure do not have any dependency between them. The only expectation is that the performance measure should be some meaningful measure of reaching a goal from a state.

Reviewer comment (Writing):
In page 5, what does it mean by ‘combinatorial generalization’? Do you have a definition or any reference for that?

Author response:
> what does it mean by ‘combinatorial generalization’? 

We have added a reference that studies combinatorial generalization in the context of RL. But the similarity is that stitching generalization tests on previously unseen combinations of states and goals which have been separately seen.

Reviewer comment (Theory):
Just curious: regarding Lemma 4.1, do you have any comments or implications on the (sort of) bias term $\E_{p(h)} \left[ p^{\beta_h}_+ (s_{t+} \mid s) p^{\beta_h}_+ (s) \right] - p^{\beta}_+ (s_{t+} \mid s) p^{\beta}_+ (s)$ ?

Author response:
> do you have any comments or implications on the bias term

In this bias term, the marginal distribution over states is the same for both sides (Lemma 3.1). But given any state, the support over the goal space of RHS will always be wider than the support over the goal space of LHS. This can be easily proved. For any given state s,
If any of the data collecting policies visits a goal g, this implies that the BC policy will also visit the goal g. By definition the BC policy will imitate all data collecting policies with some non zero probability.

But the converse of this is not true in general, that is there will be a goal which the BC policy will visit, but no data collecting policies will visit (lemma 4.1). Hence the implications of this gap are methods that do stitching and can reason / visit a wider distribution over goal from any given state when compared to methods that cannot do stitching.

Reviewer comment (Writing):
Some key concepts and results are not clearly explained.

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

Reviewer comment (Writing):
The core concept, "stitch property", is poorly introduced and readers have to refer to the reference to know what it means. There is no consistent nor clear definition for this "stitch property", making it difficult to understand what the paper is trying to achieve.

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

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

Reviewer comment (Theory):
The proposed heuristic of using clustering to stitch close states together lacks theoretical guarantees. What distance metric is suitable here? In the RL context, similar states do not mean they have similar outcomes. For example, two states can be close to each other in Euclidean norm but blocked by a walk and thus not reachable; or two states can be far apart (e.g., top and bottom of a cliff) but one can easily reach the bottom by falling but not the other way around.

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

Reviewer comment (Evaluation):
For the experiment, the paragraph on "Does more data remove the need for augmentation" is not very surprising given the construction of the dataset. In SL, the fact that more data helps generalization is based on the IID assumption, which does not hold in the current setting since the (state, goal) pairs are not seen during training (also Lemma 4.1).

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

Reviewer comment (Writing):
Both generalization and generalisation are used.

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

Reviewer comment (Writing):
point-maze large taks -> task

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

Reviewer comment (Reproducibility):
Any suggestions for a reasonable distance metric for the states?

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

Reviewer comment (Theory):
Are there any theoretical justification for the temporal augmentation?

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

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

Reviewer comment (Theory):
The proposed stitching generalization describes how "far away" different goals are in the sense of transitions, and the proposed method relies on a heuristic to describe that. This is basically what we aim to learn in TD learning and in most cases we cannot just get this kind of information for free. See Q1.

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

Reviewer comment (Experiments):
The experiments are limited to Antmaze where the proposed clustering would probably work with raw state inputs, but if we imagine pixel inputs, it would be very hard to find some heuristic to make it work, as mentioned in 1.

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

Reviewer comment (Experiments):
Missing baselines: Contrastive-based learning is also studied for goal-reaching problems as one kind of supervised learning problem, and it might also learn some representation more useful than the baseline included in the paper.

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

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.

Reviewer comment (Theory):
The optimal augmentation would be some metric that describes "how 'far away' different states are in the sense of transitions" almost perfectly. Then such a metric will basically provide us with the optimal value function. In this sense, getting such a metric could be as hard as solving it with TD, so why does the proposed method necessarily "close the gap" between TD and SL?

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

Reviewer comment (Theory):
Some of the theoretical arguments fail to convince. 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.

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

Reviewer comment (Experiments):
The experimental evaluation seems limited overall, and somewhat preliminary. I usually do not put much weight on this kind of weakness, but I also am not comfortable with the theoretical positioning of the paper. Hence, I must put more weight on the experiments in my decision for this paper.

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

Reviewer comment (Theory):
Section 1 (Stitching as goal-state pairs): My understanding of the stitching property is that it is more general: given two partially disjoint trajectories that are partially optimal, stitching via dynamic programming can combine the two partially optimal trajectories towards an optimal trajectory. Is your definition equivalent, is it a generalization or is it specific to the goal-directed setting?

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

Reviewer comment (Writing):
Section 1 (Limited data vs generalization): You state that this is not a problem of limited data, because "there can be (state, goal) pairs that are never visited in the same trajectory, despite being frequented in separate trajectories." In a way this is a source of limited data, if the "data" is a trajectory rather than an experience tuple. I think this is implicit in your argument, but it can be made explicit for clarity.

Author response:
We would like to reaffirm that this is not a source of limited data. For example in Figure 1, even after collecting infinite data from the blue and red data collecting policies, we would never see (state 2, goal 4) together. The reason why this isn’t a problem of partial observability is that you do not need the “context” information to recover the distribution of the BC policy. TD learning methods also do not have access to the context information, but can still perform stitching.

Reviewer comment (Theory):
Section 4 (More formal "Limited data vs generalization"): I am still not sure how this is different from iid generalization, or how the lemma 4.1 makes this point. My understanding of lemma 4.1 is that there exists a context distribution for which the distribution of states induced by a collection of policies will never be equal to the BC policy. But this is because the policies collecting the data are conditioned on information not available to the BC policy. This induces something like partial observability, which is the source of the problem. Without the context, I do not think this would hold and thus stitching generalization would be equivalent to iid generalization. But the problem is, in a sense, limited data: the context is never accessible and hence the data (or information) is fundamentally limited.

Author response:
Stitching is defined as a different form of generalization (not iid) because test and train distribution for stitching are different. You are correct to say that if the offline data was collected by a single data collecting policy then the train and test distribution would become the same. The case where offline data is collected by a single data collecting policy,  is covered by our framework where there is just one context with probability 1. For example, in Figure 1, if the underlying dataset was actually collected by the purple policy, we would eventually see (state 2, goal 4) together.

Reviewer comment (Theory):
Section 4 (stitching is not finding an optimal policy): I agree that stitching can help find an optimal policy, but this is besides the point in the context of your setting (offline RL). The question is two-fold, under what conditions: 1) can stitching help find a better policy and 2) does finding a better policy necessitate stitching.

Author response:
Intuitively, stitching enables OCBC algorithms to improve their performance on state-goal pairs which wouldn’t be seen otherwise. For example in Figure 1, without stitching, the OCBC policy would never be able to navigate from state 2 to state 4. To answer both your questions : for state-goal pairs that only occur across trajectories of different data collecting policies, stitching is both necessary and sufficient to find a better policy. In general, for any state-goal pairs, we can show that the optimal stitching policy is an improvement over the BC policy (See appendix D.1).

Reviewer comment (Theory):
Section 5 (Nearby states): How important is it that the states are mapped exactly onto the states in the set of experience? For example, you could imagine adding noise to all states and goals, which would therefore naturally stitch several nearby trajectories. Given the fact that you are assuming the space to have a distance metric, it is probably well-behaved so that this is effective.

Author response:
It is not important to have exactly similar states. In our ant experiments, the observations are R^29 and it is very unlikely that the same floating point number is repeated twice. Adding noise to states and goals  will not make the test and train distributions over states, goals similar. In Figure 3 (left), in the pointmaze umaze dataset, a goal in the yellow region will never be sampled from a state in the red region. Also, adding larger amounts of noise will sample invalid states (for eg walls). Additionally, our new image experiments (Fig 9 and Fig 10) suggest that our method also works well in high dimensional settings. We welcome your suggestions and are happy to conduct an experiment if you can give us a way to add noise.

Reviewer comment (Evaluation):
Section 5 (conditions for evaluating stitching): While the first condition is easy to engineer into the problem, the second condition seems problematic. How can you know whether the BC policy has a non-zero probability besides running the experiment?

Author response:
While stitching is a good property for algorithms to have [1,2, 4], we don't always want to optimize for it in isolation for real world datasets. To evaluate whether your algorithm can stitch, it is easy to create datasets satisfying both conditions. For example, in Figure 3, test on state-goal pairs which lie in different colored regions. Moreover, we have open sourced our datasets, and code to create these new such datasets for future researchers (see supplementary material).

Reviewer comment (Experiments):
Section 6 (DT + aug?): One interesting approach that is missing is combining DT with data augmentation, seeing as DT is more performant than RvS. Is there any reason why this was left out? Is it computational concerns or a more fundamental limitation.

Author response:
Because DT [5] has memory, the data augmentation method should cluster a history of states and actions, which was difficult to do computationally. In the caption of Figure 4 (green text), we have added this clarification.

Reviewer comment (Theory):
Section 6 (More data): This seems at odds with condition 1 under "popular datasets do not evaluate stitching". Surely, if every state and goal pair were included in the dataset, then DT would improve?

Author response:
As we explained above, if all state and goal pairs are already present in the dataset, then you won’t be able to test for stitching because everything will be in-distribution. The problem with the original D4RL datasets is that they evaluate on the state-goal pairs, which are visited by the same data collection policy. Hence it is possible for DT to perform well because it can capture the training distribution well. This is unlike our datasets, where there is a clear difference between methods that perform stitching (Data augmentation) and methods that don’t (OCBC).