Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View.
Authors: Raj Ghugare
Matthieu Geist
Glen Berseth
Benjamin Eysenbach
Venue: ICLR 2024 poster (2024)
Keywords: reinforcement learning
data augmentation
stitching
Primary area: reinforcement learning
TL;DR: We show that the stitching property in RL is a form of generalization that SL-based methods can not have, however we show how a new type of data augmentation facilitates this type of generalization.

Abstract: Some reinforcement learning (RL) algorithms have the capability of recombining together pieces of previously seen experience to solve a task never seen before during training. This oft-sought property is one of the few ways in which dynamic programming based RL algorithms are considered different from supervised learning (SL) based RL algorithms. Yet, recent RL methods based on off-the-shelf SL algorithms achieve excellent results without an explicit mechanism for stitching; it remains unclear whether those methods forgo this important stitching property. This paper studies this question in the setting of goal-reaching problems. We show that the desirable stitching property corresponds to a form of generalization: after training on a distribution of (state, goal) pairs, one would like to evaluate on (state, goal) pairs not seen together in the training data. Our analysis shows that this sort of generalization is different from i.i.d. generalization. This connection between stitching and generalization reveals why we should not expect existing RL methods based on SL to perform stitching, even in the limit of large datasets and models. We experimentally validate this result on carefully constructed datasets.
This connection suggests a simple remedy, the same remedy for improving generalization in supervised learning: data augmentation. We propose a naive temporal data augmentation approach and demonstrate that adding it to RL methods based on SL enables them to successfully stitch together experience, so that they succeed in navigating between states and goals unseen together during training.

CLOSING THE GAP BETWEEN TD LEARNING AND SUPERVISED LEARNING - A GENERALISATION POINT OF VIEW
Raj Ghugare'
Matthieu Geist
Glen Berseth
Benjamin Eysenbach 3+
Mila, Université de Montréal
2 Google DeepMind
SPrinceton University
aag.ghugarcemiie www.ggrrrremm....eeee.eee..ee
ABSTRACT
Some reinforcement learning (RL) algorithms can stitch pieces of experience to solve task never seen before during training. This oft-sought property is differ
of the few ways in which RL methods based on (SL). Yet, certain RL methods
from RL methods based on supervised-leaniin based on off-the-shelf SI. algorithms achieve excellent results without an explicit mechanism for stitching; it remains unclear whether those methods forgo this important stitching property. This paper studies this question for the problems of achieving target goal state and achieving a target return value. Our main result is to show that the stitching property corresponds to form of combinatorial generalization: after training on a distribution of (state, goal) pairs, one would like to evaluate on (state, goal) pairs not seen together in the training data. Our analysis shows that this sort generalization is different from i.i.d. generalization. This connection between stitching and generalisation reveals why we should not expect SL-based RL methods to perform stitching, even in the limit of large datasets and models. Based this analysis, we construct new datasets to explicitly test for this property, revealing that SL-based methods lack this stitching property and hence fail to perform combinatorial generalization. Nonetheless, the connection between stitching and combinatorial generalisation also suggests simple remedy for improving generalisation in SL: data augmentation. We propose a temporal data augmentation and demonstrate that adding it to SL-based methods enables them to successfully complete tasks not seen together during training. On a high level, this connection illustrates the importance of combinatorial generalization for data efficiency in time-series data beyond tasks beyond RL., like audio, video, or text.
INTRODUCTION
Many recent methods view RL as purely SL problem of mapping input states and desired goals, to optimal actions [.3. These methods have gained lot of attention due to their simplicity and scalability 4. These methods sample goal g (or return r) from the dataset, which was previously encountered after taking an action a from a state 8, and then imitate a by treating it as an optimal label for reaching y (or achieving return r) from 8. These methods, collectively known as outcome conditional behavioral cloning algorithms (OCBC), achieve excellent results on common benchmarks 3. However, a fundamental level, there some important differences between RL and SL. This paper studies one of those differences: the capability of some RL algorithms to stitch together pieces of experience to solve task never seen during training. While some papers have claimed that some OCBC approaches already have this stitching property (2), both our theoretical and empirical analyses suggest some important limitations of these prior claims.
The stitching property 55] is common among RL algorithms that perform dynamic programming DQN
8 IQL (). It is often credited for multiple properties of dynamic programming algorithms like superior data efficiency and off policy reasoning (See Section 4 for detailed discussion). Importantly, we show that stitching also allows for third property the ability to infer solutions to combinatorial number of tasks during test time, like navigating between certain state-gon goal pairs that never appear together (but do appear separately) during training. example of stitching is that humans don't 't need access to optimal actions to go from an airport new tourist places; they can use their previous knowledge to navigate to a taxi-stand, which would take them to any
(eN (6. DDPG (7), TD3
"Equal advising.

---

location. But, purely supervised approaches to sequential problems like RL, do not explicitly take such temporal relationships into account. Even in other sequential domains like language, a large body of work is dedicated to study the combinatorial generalisation abilities of large language models [3] Our work shows that combinatorial generalisation is also required to solve tasks in the context of RI..
We start by formalising stitching as form of combinatorial generalisation. We observe that when data are collected from a mixture of policies, there can be certain (state, goal) pairs that are never visited in the same trajectory, despite being frequented in separate trajectories. Information from multiple trajectories should be stitched to complete these tasks. Because such tasks (state-goal pairs) are seen separately, but never together, we call the ability of algorithms to perform these tasks as combinatorial generalisation. This connection further motivates an inspiration from SL; if generalisation is the problem, then data augmentation likely an effective approach [14]. We propose a form of temporal data augmentation for OCBC methods so that they acquire this stitching property and succeed in navigating between unseen (start, goal) pairs or achieving greater returns than the offline dataset. This form of data augmentation involves fime, rather than random cropping or shifting colors. Intuitively, temporal data augmentation augments the original goal, with a new goal sampled from different overlapping trajectory in the offline dataset. This data augmentation scheme does require an estimate of distance between states to detect overlapping trajectories. We demonstrate that this data augmentation is theoretically backed, and empirically endows OCBC algorithms with the stitching property on difficult state-based and image-based tasks.
Our primary contribution is to provide formal framework for studying stitching as form of combinatorial generalisation. Because of this connection, we hypothesize that OCBC methods do not perform stitching. Perhaps surprisingly, simply increasing the volume of data does not guarantee this sort of combinatorial generalization. Our empirical results support the theory: we demonstrate that prior RL methods based on SL (DT 2 and RvS BD) fail to perform stitching, even when trained on abundant quantities of data. Our experiments reveal a subtle consideration with the common D4RL, datasets [5]: while these datasets are purported to test for exactly this sort of combinatorial generalization, data analysis reveals that "unseen" (state, goal) pairs do actually appear in the dataset. Thus, our experiments are run on a new variant of these datasets that we constructed for this paper to explicitly test for combinatorial generalization On 10 different environments, including both state and image based tasks, and goal and return conditioning, adding data augmentation improves the generalisation capabilities of SL approaches by up factor of 2.5.
RELATED WORK
Prior methods that do some form of explicit stitching. Previous work on stitching abilities of SL algorithms have conflicting claims. The DT paper [2) shows experiments where their SL-based method performs stitching. On the contrary. [66] provide an example where SL algorithms do not perform stitching. RyS [3] shows that a simple SL-based algorithm can surpass the performance of TD algorithms. In tabular settings, 17 show that the benefits of TD-leamnng arise from trajectory stitching. We provide a formal definition of stitching contrast, generalisation in RL has been generally associated with making correct predictions for unseen but similar states and actions (1) planning [21], ignoring irrelevant details [22-24)) or robustness towards changes in the reward or transition dynamics 27
form of combinatorial generalisation. In
Offiine RL datasets. A large amount of work done to build offline RL datasets. 15 provided first standard offline RI benchmark, (8 provide exploratory offline datasets to underscore the importance of diverse data, (29 30 focus on data efficiency and real world deployement and 31 provide benchmarks that also compare the online evaluation budget of offline RL algorithms. Although many offline RL papers informally allude to stitching, we devise new offline RL datasets that precisely test the stitching abilities of offline RL algorithms.
Data augmentation in RL. Data augmentation has been proposed as remedy to improve generalisation in RL (32, akin to SL (39]. Perhaps the most similar prior work are the ones which use dynamic programming to augment existing trajectories to improve the performance properties of SL algorithms 40 42 However, because these methods still require dynamic programming, they don't it have the same simplicity that make SL algorithms appealing in the first place.
'Open sourced code and data is available: stitching- www.nngi-isiiiiiiaaaaeaaeaeiii combi -genrral saaioon
https://github. com/ Ra jGhugarel
2

---

3 PRELIMINARIES
Controlled Markov processes. We will study the problem of goal-conditioned RL in a controlled Markov process with states E S and actions a E A. The dynamics are p(s' s, the initial state distribution is po (so), the discount factor is 7. The policy π T(a, 8, g) is conditioned on a pair of state and goal s.g S. For policy म, define Pp (st so) as the distribution over states visited after exactly steps. We can then define the discounted state occupancy distribution and its conditional counterpart as
p (s++ =g) E Es-po(so) [5+ = | = =s]]
(1)
p- (SL+ = 9 su = 5) - (1 Σ Y'P = | = s).
(2)
where 924 is the variable that specifies a future state corresponding to the discounted state occupancy distribution. Given a state-goal pair s; ~> (ses(,, test time, the task of the policy is to maximise the probability of reaching the goal in the future
max (T); where [(T) (e.gr-paa(ss9)
[ (8++ = 9 80 =)]]
Data collection. Our work focuses the offline RL setting where the agent has access fixed dataset of N trajectories D = (fs), ag. )), Our theoretical analysis will assume that the dataset collected by set of policies {B(a s, h)) where h specifies some context. For example. could reflect different goals, different language instructions, different users or even different start state distributions. Precisely, we assume that the data was collected by first sampling context from a distribution p(h), and then sampling trajectory from the corresponding policy B(a 8, h). We will use the shorthand notation 1.)) = 3- :, h) to denote the data collecting policy conditioned on context h. Trajectories are assumed to be stored without h, hence the context denotes all hidden information that the true data collection policies used to collect the data.
This setup of collecting data corresponds to mixture of Markovian policies There is a classic result saying that, for every such mixture of Markovian policies, there exists a Markovian policy that has the same discounted state occupancy measure.
Lemma 3.1 (Rephrased from Theorem 2.8 of 43), Theorem 6.1 of (44). Let a set of context- conditioned policies (Ar(a 8)} and distribution over contexts p(h) be given. There exists Markovian policy B(a s) such that it has the same discounted state occupancy measure as the mixture of policies:
(+(+) = Ep(h)
8n p (s(+)
(4)
The policy (a s) is simple to construct mathematically as follows. For data collected from the mixture of context conditioned policies, let p² (Ne s) be the distribution over the context given that the policy arrived in state 8.
8(a Σ (a s)p (11 h
Theorem 6.1 44 proves the correctness of this construction. The policy B(a 8) is also easy to construct empirically simply perform behavioral cloning (BC) on data aggregated from the set of policies. We will hence call this policy the BC policy.
Outcome Conditional behavioral cloning (OCBC). While our theoretical analysis will consider generalisation abstracted away from any particular RL algorithm, we will present empirical results using a simple and popular class of conditond RL methods: Outcome conditional behavioral cloning (45) (DT (2). URL M1. RvS 3 GCSL (46) and many others 47 48) ). These methods take
input dataset of trajectories D ((+6) ab DA and learn a goal-conditioned policy л<a 8, g)
using a maximum likelihood objective:
max E (e (logT(a) 8.
(6)
*(+++)
Note that the mixture is the level of trajectories, at the level of individual actions.
3

---

The sampling above can be done by first sampling a trajectory from the dataset (uniformly random), then sampling a (state, action) pair from that trajectory, and setting the goal to be a random state that occurred later in that same trajectory. If we incorporate our data collecting assumptions, then this sampling can be written as
max E En-p(()) m()b:))
E an (alo), x (e) setmuph (se+1s;a)
[log TTa 8, 8,,{)]]]-
"STITCHING" AS A FORM OF COMBINATORIAL GENERALISATION
Before concretely defining stitching, we will describe three desirable properties that colloquially associated with "stitching" and the learning dynamics of TD methods. (Property 1) The ability to select infrequently seen paths that are more optimal than frequent ones. While shorter trajectory between the state and the goal may occur infrequently in the dataset, TD methods can find more exam- ples of this trajectory by recombining pieces of different trajectories. thanks dynamic programming. This property is enjoyed by both SARSA (expectation over actions) and Q-learning (max over actions) methods, and primarily associated with the sample efficiency of learning. (Property 2) The ability to evaluate policies different from those which collected the data, and perform multiple steps of policy improvement. This property is unique to Q-learning. (Property 3) Temporal difference methods (both Q-learning and SARSA) can also recombine trajectories to find paths between states never seen together during training. This property different from the first property in that it is not a matter of data efficiency temporal difference methods can find paths that will never be sampled from the data collecting policies, even if given infinite samples. All three of these properties are colloquially referred as "stitching" in the literature. While these properties are not entirely orthogonal, they distinct: certain algorithms may have just some of these properties. Obtaining all these properties in a simpler (than TD) framework difficult, and remains unclear whether OCBC methods possess any of them. To better understand the differences and similarities this study focuses on the third property. We formalize this property as a form of generalisation, which we will refer to as combinatorial generalisation.
Defining combinatorial generalisation will allow us to analyze if and when OCBC methods perform stitching, both theoretically (this section) and experimentally (Section Intuitively, combinatorial generalisation looks at connecting states and goals, which are never seen together in the same trajectory, but where a path between them is possible using the information present in different trajectories. It therefore tests a form of "stitching" [I5] 44]. akin to "combinatoral generalisation" (50]- [2] To define this generalisation, we will specify a training distribution and testing distribution. The training distribution corresponds to sampling a context h ~ p(h) and then sampling an (s, g) pair from the corresponding policy 8r. This is exactly how OCBC methods The testing distribution corresponds to sampling an (s,g) pair from the BC policy Baa |) defined in Equation 55) For each distribution, we will measure the performance fatt--19)( (s,g) .) of goal- conditioned policy (a | 8,
trained in practice (Section B).
Definition 1 (Combinatorial generalisation). Let a set of context-conditioned policies (B s)
be given, along with a prior over contexts p(h). Let B(a s) be the policy constructed via Eq. (5)
Let Taa) 9, be policy for evaluation. The combinatorial generalisation of policy T(a s, g)
measures the differences in goal-reaching performance for goals sampled g ~ p. (su+ () versus
goals sampled from ~ Ep(h)pp (st+ s)];
(s.)
E R amp (6) (mtl1).0)s.g))
E (
(8)
a- (811 taain performance
test performance
The precise way performance is measured is not important for our analysis: "generalisation" simply means that the performance under one distribution is similar to the performance under another. In our experiments, will look at performance measured by the success rate at reaching the commanded goal. On the surface, it could seem like both the test and train distributions are the same. Lemma [3.] about reducing mixtures of policies to single Markovian policy seems to hint that this might be true.
4

---

Indeed, this distinction has not been made before while analysing OCBC methods [6 45 misconception is demonstrated by the following lemma:
This
5
50%
3
3
3
4
2
2
2
(a) MDP.
(b) 8m=1, Bh=
(c) B.
Figure I: (a) The MDP has states and two actions (up and right). (b) Training distribution: Data collected using two contexts conditioned policies shown in blue and red. (c) Testing distribution: The behavior cloned policy (equation is shown in purple. During training. the state-goal pair {s. = 2, Sct is never sampled, no data collecting policy goes from state to 4. But the behavior cloned policy has zero prohability of sampling the state-goal pair {st = 2, 8t- = 4}. Because this discrepancy between the train and test distributions, OCBC algorithms do not have any guarantees of outputting the correct action for the state-goal pair {st = 2. s1t 4). Whereas dynamic programming based methods can propagate rewards through the backwards stitched path of 4 + - 2 output the correct action.
Lemma 4.1. There exist a collection of policies {8r} and context distribution p(h) such that, conditioned on state, the distribution of states and goals for the data collecting policies (training) is different the distribution of states and goals (testing) for BC policy 3.
Ep(i)) p (80+ 8))f (t) # vL ()) for some states 8, 8t4+
Proof. The proof base simple counterexample, shown in Fig. 1 See the related caption for sketch of proof and Appendix D. for the formal one.
In summary, while the BC policy 3(a s) will visit the same states as the mixture of data collecting policies average, conditioned distribution of future states than the mixture of policies. Even infinite data is collected from the data collecting policies, there can be pairs of states that will never be visited by any one data collecting policy in a single trajectory. The important implication of this negative result is that stitching requires the OCBC algorithm to recover distribution over state-goal pairs (B) which is different from the one it is trained on (3m=1, B1=2).
some state, the BC policy 3(a 8) may visit a different
In theory, the training distribution has enough information to recover the state-goal goal distribution of the BC policy without the knowledge of the contexts of the data collecting policies. It is upto the algorithm extract this information. Many RL methods can recover the test distribution implicitly and sidestep this negative result by doing dynamic programming (i.e., temporal difference learning). One way of viewing dynamic programming that it considers all possible ways of stitching together trajectories, and selects the best among these stitched trajectories. But OCBC algorithms based on
SL [3 46 47
can only have guarantees for iid generalisation (53) And in line with previous works
54 it is not clear apriori why these studying other forms of combinatorial generalisation in SL [10
methods should have the combinatorial generalisation property, leading to the following hypothesis:
Conditional imitation learning methods do not have the combinatorial generalisation property.
We will test this hypothesis empirically in our experiments. In Appendix A we discuss connections between stitching and spurious correlations.
5 TEMPORAL AUGMENTATION FACILITATES GENERALISATION
The previous section allows us to rethink the oft-sought "stitching" property as a form of generalisation, and measure that generalisation in the way we measure generalisation in
5

---

SL: by measuring difference in performance under two objectives. Casting stitching form of generalisation allows us to employ a standard tool from SL: data augmentation. When the computer vision expert wants model that generalize to random crops, they train their model on randomly-cropped images. Indeed, prior work has applied data augmentation to RI. to achieve various notions of generalisation 32 [52] However, we use different type of data augmentation to facilitate stitching. In this section, we describe a data augmentation approach that allows OCBC methods improve their stitching capabilities.
Recall that OCBC policies are trained on (s, a. g) triplets. To perform data augmentation, we will replace 9 with different goal ğ. To sam- ple these new goals ÿ. we first take the original goal and identify states from the offline dataset which are nearby to this goal ( (Section [5). Let denote one of these nearby "waypoint" states. Looking at the trajectory that contains w, the new goal is random state that occurs after w in this trajectory. We visualize this data augmen- tation in Fig.
g
8
9
W
DATA AUGMENTATION FOR STITCH- ING: After sampling initial training example (s, a. g) (Eq. 7.. we look for a waypeint state w in the light blue region around the original goal g. and then sample a new augmented goal ğ from later in that trajec- tory. This is simple approach sample cross trajectory goals ğ such that the action still an optimal action state s.
Figure 2:
Identifying nearby states. The problem of sampling nearby states can be solved by clustering all the states from the offline dataset before training. This assumes a distance metric in the state space. Using this distance metric, every state can be assigned a discrete label from l different categories using clustering
algorithm [55] [56.. Although finding a good distance metric is difficult in high-dimensional settings [57] our experiments show that using simple L2 distance leads to significant improvement, even for image-based tasks.
Method summary. Algorithm I summarizes how our data augmentation can be applicd on top of existing OCBC algorithms. Given method to group states, we can add our data-augmentation to existing OCBC algorithms about 5 lines of code (marked in blue). In our experiments, we use the k-means algorithm [55]. To sample nearby waypoint states, randomly sample state from the same group (same cluster) as the original goal. The augmented goal is then sampled from the future of this waypoint (See Fig.
Algorithm 1 Outcomee-conditioned behavioral cloning with (temporal) data augmentation. The key contribution of our paper is this form of data augmentation, which is shown in blue text.
1: Input: Dataset D ({ so. do.... })-
Initialize OCBC policy no [alss,9) 9) with parameters 0.
3: Set e augmentation probability, m = mini--btch size.
4: ((do, di.... 1)= CLUSTER( Sp, $1.... D).
Group all states in the datasct.
5: while not converged do
6: for 1.... - do
Sample (st, dr, 9c+) ~ D. 8: With probability e:
7:
Equation
9:
Get the group of the goal: k = de+-
10:
Sample waypoint states from the same group: w {s.: Vi such that d, k}) Sample augmented goal ğ from the future of w. from the same trajectory as w.
11:
12:
Augment the goal G++ = ğ-
13:
Collect the loss La (f) = log (a (t,0)))-
14:
Update 0 using gradient descent on the mini-batch loss 1 m
Σ (.)
15:
Return: Ta(a)s.g)
Theoretical intuition on temporal data augmentation. While data augmentations in general do not have exact theoretical guarantees, we can prove that temporal data augmentation, under certain

---

smoothness assumptions, will generate additional stateggoal pairs which may not be otherwise during training. Appendix D22 we show that there exists a hierarchy increasing stitching abilities, where O-step distribution corresponds to the train distribution (Eq. (9).
distributions with
left) and the per-step distribution corresponds to the test distribution (Eq. (9) right). We prove that applying temporal data augmentation once, samples state-goals from the 1-step distribution.
Lemma 5.1. Under the smoothness assumptions mentioned in Appendix D.2 (Eq. (T)). for all 5, a pairs, temporal data augmentation poomp aug (g s, a) approximately samples goal according the distribution one-step stitching policy (pl-ser) (g 8, a.).)
Intuitively, the smoothness assumptions are required to ensure that nearby states have similar proba- bilities under the data collection policies. In Fig. f example, this ensures taking action from state s has similar probabilities of reaching nearby states w and For the complete proof as well as more details see Appendix D.2
SUE
Umaze
Medium
Large
Miniworld
Figure 3: Goal conditioned RI Different colors represent the navigation regions different data collecting policies. During data collection, these policies navigate between random state-goal pairs chosen from their region navigation. These visualisations are for the "point" mazes. The "ant" maze datasets are similar. Appendix Fig. 12 shows the "ant" maze datasets.
Return condi- tioned RI. We visualise our new image based and partially observable environment cre- ated using Miniworld (58).
Figure 4:
EXPERIMENTS
The experiments aim (1) to verify our theoretical claim that OCBC methods do not always exhibit combinatorial generalisation, even with larger datasets larger transformer models, and (2) to evaluate how adding temporal augmentation to OCBC methods can improve stitching in both state- based and image-based tasks. All experiments are conducted across five random seeds.
OCBC methods. RvS 3 is an OCBC algorithm that uses fully connected policy network and often achieves results better than TD-learnnn algorithms on various offline RL benchmarks (3]. DT 2 treats RL as sequential SL problem and uses the transformer architecture as policy. DT outputs an action, conditioning not only on the current state, but a history states, actions and goals. See Appendix C.2 for implementation details.
6.1 TESTING THE ABILITY OCBC ALGORITHMS AND TEMPORAL DATA AUGMENTATION TO PERFORM STITCHING.
While the maze datasets from D4RL 15] were originally motivated test the stitching capabilities of RL algorithms, we find that most test state-goal pairs are already in the training distribution. Thus, good success rate on these datasets does not necessarily require stitching. This may explain why OCBC methods have achieved excellent results on these tasks B3, despite the fact that our theory suggests that these methods do not perform stitching. In our experiments, we collect new offline datasets that precisely test for stitching (see Fig. 3 and Fig. 12 for visualisation). To collect our datasets, we use the same "point" and "ant" mazes (umaze, medium and large) from D4RL (15). To test for stitching, we condition OCBC policies to navigate between (state, goal) pairs previously unseen together, and measure the success rate. In Fig. this conditioning corresponds to (state, goal) pairs that appear in differently coloured regions. Each task consists of 2-6 randomly chosen (state, goal) pairs from different regions the maze. In Appendix C33 we discuss the important differences between the D4RL and our datasets, which necessary to test for stitching.
7

---

Rv5
RvS state augmentation (ours)
RvS only augieentttion (ours)
DT
DT + state augmantation (ours)
DT + only goal augmentation (ours)
1.03
3.500
1.0
Success rate
0.75
1375
0.75
0.375
0.53
1252
0.50
0.253
0.23
2123
umaze
medium
large
umaze
medium
large
umaze
medium
large
umaze
medium
large
Figure Adding data augmentation outperforms the OCBC baselines on most tasks. "Only goal augmentation" refers to an oracle version of our augmentation that uses privileged information (x; y coordinates) when performing augmentation. Adding temporal data augmentation (both standard and oracle versions) improves the performance of both RvS and DT on 8/a tasks.
Results. In Fig- we can that both DT and RvS struggle to solve unseen tasks test-time. However, applying temporal data-augmentation to RvS improves the goal-reaching success rate on 6/6 tasks, because the augmentation results in sampling (state, goal) pairs otherwise unseen together. To show that temporal data augmentation can also be applied to only important parts of the state, based on extra domain knowledge, we also compare an oracle version of our data augmentation. This oracle version uses only the coordinates from the state vector to apply the K-means algorithm. Figure 5) also shows that using extra domain knowledge can further improve performance.
RvS
RvS data augmen
DT Data Augmentation
visual polnt-maze
visual point-maze
miniworid
1.00
2.0
0.75
surcens rate
3.5
Episedic
0.25
0.25
medium
large
medium
large
collect
collect
Figure 6: Temporal data augmentation image-based tasks. It is difficult to find reliable metric apply temporal data augmentation high- -igeenimnni tasks, We show that using simple L2 distance metric can surprisingly improve the combinatorial generalisation of OCBC algorithms on both goal-conditioned (left and center) and return-conditooned (right) tasks.
6.2 CAN TEMPORAL DATA AUGMENTATION WORK FOR HIGH DIMENSIONAL TASKS?
As mentioned in Section 5 it can be difficult provide good distance metric, especially for tasks with nigh-dimensional states. Although this a limitation, we show that temporal data augmentation, using simple L2 distance metric, can improve the combinatorial generalization OCBC algorithms even on high-dimensione image-based based tasks. To evaluate this, we use both image-basee based goal- conditioned and return-come coturnccondiioonedd tasks. For the goal-conditioned tasks, we use an image-based version of the "point" mazes 3 The agent is given top-down view of the maze to infer its location. For the return conditioned tasks, we create new task using Miniworld (58) (See Fig. 3 called "collect". The task is to collect both the keys and return to the start position. A reward of is received after collecting each key. There are two data collecting policies, each collecting only one key. test time, the OCBC policy is conditioned the unseen return of (collect both keys).
Results. In Fig. 6 we can see that temporal data augmentation improves the performance of RvS and DT 4/4 and 3/4 tasks, respectively. Although temporal data augmentation can be successfully applied to some high dimensional tasks, is not guaranteed to succeed There remains room for other scalable and robust methods to achieve even better performance.
8

---

(ours)
DT a data augmentatin (ours)
0.75
0.75
0.50
0.50
Ied dete dos, hidden layers
Ied stow, hidden legers
saccess mete
0.50
250
SUCCCSE rate
0.00
5.00
105
1e6
167
0.00
0.00
dataset
number hidden layers
Figure 7: Performance of DT trained different offline dataset sizes (left) and using a different num- ber of hidden layers (right) averaged across all "point" mazes. Even with larger datasets models, the gener- alisation of DT is worse than DT + data augmentation.
Figure 8: Comparison of data augmentation trained with different numbers of centroids the K-mcans algorithm on "point" maze-medium and "ant" maze-medium. Temporal data augmentation corresponding to all values of l outperforms DT both tasks.
6.3 ABLATION EXPERIMENTS.
Does more data remove the need for augmentation? Although our theory (Lemma 4.1) suggests that generalisation is required because of a change in distribution and is not problem due to limited data, conventional wisdom says that larger datasets generally result in better generalisation. To empirically test whether this is the case, we train DT on 10 million transitions (10 times more than Fig. () on all "point" maze tasks. In Fig. (left), we see that even with more data, the combinatorial generalisation of DT does not improve much. Lastly, scaling the size of transformer models [59] is known to perform better in many SL problems. To understand whether this can have an effect on stitching capabilities, increased the number of layers in the original DT model. In Fig. (right), we can see that increasing the number of layers does not have an effect on DT's stitching capabilities.
How sensitive temporal data augmentation the number of centroids used for K-means? In Fig. [ we ablate the choice of the number of centroids used in K-means two environments "point" maze-men maaze-medium and "ant" maze-medium All choices centroids significantly outperform the RyS method on both tasks.
Combinatorial generalisation due to spurious relations. In most of our experiments, OCBC algorithms do exhibit, albeit very low, combinatorial generalisation. We believe this occurs not due to the combinatorial generalisation ability of OCBC algorithms, but due certain spurious relations that are present in the dataset. In Appendix(A) [ we discuss the relation of combinatorial generalisation with spurious relations. In Appendix B perform didactic experiments to show that combinatorial generalisation in OCBC algorithms occurs because the OCBC policy network picks up on such spurious relations.
DISCUSSION
In this work, we shed light on an area that the community has been investigating recently, can SL-based approaches perform stitching. We formally show that stitching requires combinatorial generalisation, and recent SL. approaches to RL (OCBC methods) generally not have any guarantees to perform such generalisation. We empirically verify this many state-based and image-based tasks. We also propose a type of temporal data augmentation to perform the desired type of combinatorial generalisation precisely and help bridge the gap between OCBC and temporal difference algorithms.
Limitations. Our proposed augmentation assumes access to a local distance metric in the state space, which can be difficult to obtain in general. Lifting this assumption and developing scalable OCBC algorithms that generalise is promising direction for future work.
Overall, our work hints that current SL approaches may not efficiently use sequential data found in RL: even when trained on vast quantities of data, these approaches do not perform combinatorial generalisation (stitching). Due to the temporal nature of RL, it is possible to solve a combinatorial number of tasks from the same sequential data. Similar gains in data efficiency be made by designing algorithms capable of combinatorial generalisation in other problems involving time series data, for example, audio, videos, and text.

---

Acknowledgements. This work supported by Mila IDT, Compute Canada, and CIFAR. We thank Artem Zholus, Amav Jain, and Tianwei Ni for reviewing an earlier draft of our paper. We thank Scohong Park and Mikail Khona for helpful pointers related to the code. We thank Siddarth Venkatraman and members of the Robotics and Embodied Al (REAL) Lab for fruitful discussions throughout the project.
REFERENCES
[1] Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards just map them actions, 2020.
[2] 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 neural information prucessing systems, 34:15084-15097, 2021.
[3] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rys: What is essential for offline
rl via supervised learning? arXiv preprint arXiv:21 12. 10751, 2021.
[4] Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Damiel Freeman, Winnic Xu, Sergio Guadarrama, lan Fischer, Fric Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers, 2022.
[5] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pages 1433-1438 Chicago, IL, USA, 2008
[6] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, arXiv preprint arXiv:
2013.
[7] Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Frez, Yuval Tassa, David Silver. and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint www.:550.00999 2015.
[8] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages 1587-1596. PMLR, 2018.
[9] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- g-learning. arXiv preprint arXiv:2110.06/11 2021.
[10] Thaddaus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Compositional generalization from first principles. 2023.
[11] Abulhair Saparov and He He. Language models are greedy reasoners: A systematic formal analysis of chain-of-thought, 2023.
[12] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task, 2023.
[13] Maxwell Nye. Anders Johan Andreassen, Guy Cur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021.
[14] Luis Perez and Jason Wang. The effectiveness of data augmentation in image classification using deep learning, 2017.
[15] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2021.
[16] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does cnturn-coniitioned supervised learning work for offline reinforcement learning? Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[17] David Cheikhi and Danicl Russo. On the statistical benefits of temporal difference learning. 2023.
[18] Amy Zhang. Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. 2018
[19] Karl Cobbe, Oleg Klimov, Chris Hesse, Tachoon Kim. and John Schulman. Quantifying generalization in reinforcement learning. 2019.
[20] Kenny Young, Aditya Ramesh, Louis Kirsch, and Jürgen Schmidhuber. The benefits of model- generalization in reinforcement learning, 2023.
[21] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 1273-1286. Curran Associates, Inc., 2021.
10

---

[22]
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction, 2021.
[23] Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information prioritization through empowerment in visual model- rl, 2022.
[24]
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization in reinforcement lcarning with selective noise injection and information bottleneck, 2019.
[25] Jun Morimoto and Kenji Doya. Robust reinforcement learning. In T. Leen, T. Dietterich, and V. Tresp, editors, Advances in Neural Information Processing Systems, volume 13. MIT Press, 2000.
[26] Chen Tessler, Yonathan Efroni, and Shie Mannor. Action robust reinforeement learning and applications in continuous control, 2019.
[27]
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. Robust predictable control, 2021.
[28] Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't t change the algorithm, change the data: Exploratory data for offline reinforcement learning, 2022.
[29] Gaoyue Zhou, Liyiming Ke, Siddharthha Srinivasa, Abhinav Gupta, Aravind Rajeswaran, and Vikash Kumar. Real world offline reinforcement learning with realistic data source, 2022.
[30] Zheqing Zhu, Rodrigo de Salvo Braz, Jalaj Bhandari, Daniel Jiang. Yi Wan, Yonathan Efroni, Liyuan Wang. Ruiyang Xu, Hongbo Guo, Alex Nikulkov, Dmytro Korenkevych, Urun Dogan, Frank Cheng, Zheng Wu, and Wanqiao Xu. Pearl: A production-read reinforcement learning agent, 2023.
[31] Vladislav Kurenkov and Sergey Kolesnikov. Showing your offline reinforcement learning work: Online
evaluation budget matters, 2022.
[32] Austin Stone. Oscar Ramirez. Kurt Konolige, and Rico Jonschkowski. The distracting control suite challenging benchmark for reinforcement learning from pixels, 2021.
[33] Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi talti-task robotic reinforcement learning at scale, 2021.
[34] Aravind Srinivas, Michael Laskin, and Picter Abbeel. Curl: Contrastive unsupervised representations for reinforcement lcarning, 2020.
[35] Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation, 2021.
[36] Ila Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels, 2021.
[37] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning, 2021.
[38] Chaochao Lu, Biwei Huang, Ke Wang. José Miguel Hernández-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual-hasee data augmentation, 2020.
[39] Connor Shorten and Taghi M. Khoshgoftaar. A survey image data augmentation for deep learning. Journal of Big Data, 6:1-48. 2019.
[40] Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodrigee Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline RL, 2023.
[41] Keiran Paster. Silviu Pitis, Sheila A. Mcllraith, and Jimmy Ba. Return augmentation gives supervised RL temporal compositionality, 2023.
[42] Ian Char. Viraj Mehta, Adam Villaflor, John M. Dolan, and Jeff Schneider. Bats: Best action trajectory stitching, 2022.
[43] Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mellon University, 2010.
[44] Eugene A Feinberg and Adam Shwartz. Handbook of Markov decision processes: methods and applications. volume 40. Springer Science & Business Media, 2012.
[45] Benjamin Eysenbach, Soumith Udatha, Sergey Levine, and Ruslan Salakhutdinov. Imitating past successes can be very suboptimal. arXiv preprint arXiv:2206.00 033378, 2022.
[46] Dibya Ghosh, Abhishck Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. arXiv preprint arXiv: 2019.
[47] Hao Sun, Zhizhong Li, Xiaotong Liu, Bolei Zhou, and Dahua Lin. Policy continuation with hindsight inverse dynamics. Advances in Neural Information Processing Systems, 32, 2019.

---

[48] Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-conditioned policies, 2019.
[49] Michael Janner, Qiyang Li, and Sergey Levine. Offiie reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-122 2021.
[50] Ivan Vankov and Jeffrey Bowers. Training neural networks to encode symbols enables combinatorial generalization, 2019.
[51] Philippe JlansnnEEstruee Amy Zhang, Ashvin Nair, Patrick Yin, and Sergey Levine. Bisimulation makes analogies goal-conditioned reinforcement learning, 2022.
[52] Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76:201-264, jan 2023.
[53] Shai Shalev-Shwartz and Shai Ben-David, Understanding Machine Learning- Cambridge University Press, 2014.
From Theory to Algorithums.
[54] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains, Mach. Learn., 19(1-2): may 2010.
[55] S. Lloyd. Least squares quantization in pem. IEFF Transactions on Information Theory, 28(2): 129-137, 1982.
[56] Alexander A. Sherstov and Peter Stone. Function approximation via tile coding: Automating parameter choice. In Jean-Daniel Zucker and Lorenza Saitta, editors, Abxtraction, Reformulation and Appreximation. pages 194-205, Berlin, Heidelberg, 2005. Springer Berlin Heidelberg.
[57] Charu Aggarwal, Alexander Hinneburg. and Daniel Keim. On the surprising behavior of distance metric in nigh-dimensional space. First publ. in: Database theory, ICDT 200, 8th International Conference, London, UK, January 4. 6, 2001 / Jan Van den Bussche ... (eds.). Berlin: Springer, 2001, pp. 420-434 (=Lecture notes in computer science; ; 1973), 02 2002.
[58] Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
[59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[60] Ishita Dasgupta, Erin Grant, and Thomas L. Griffiths, Distinguishing rule- and exemplar-basse based generaliza- tion in learning systems, 2022.
[61] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arX17:1911.08731. 2019.
[62] Mark Towers, Jordan K. Terry, Ariel Kwiatkowski, John U. Balis, Gianluca de Cola, Tristan Delcu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Vicente, Andrea Pierré, Sander Schulhoff, Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymnasium, March 2023.
[63] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg. J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-lear:: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830, 2011.
[64] K. Hinderer. Lipschitz contimuity value functions in markovian decision processes. Mathematical
Methods of Operations Research, 62:3-22, 09 2005.
[65] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning, 2019.
[66] Emmanuel Rachelson and Michail Lagoudakis. On the locality of action domination in sequential decision making. In Advances in Neural Information Processing Systems, 01 2010.
[67] Alfred Müller. How does the value function of a markov decision process depend on the transition probabilities? Math. Oper. Res., 22(4))872-88 885, nov 1997.
12

---

A RELATIONSHIP WITH SPURIOUS CORRELATIONS.
Handling stitching is somewhat akin to handling spuri- ous correlations studied the SL community. In the RL setting, we want to navigate from A to B given dataset that contains some trajectories with and some with B but none with both A and B. This is somewhat analogous to common settings in object detection in computer vi- sion, where the object background is highly indicative of the object class. For example, the common waterbirds dataset [6]] aims to classify images of birds into two classes, "land birds" and "water birds," but the image backgrounds correlated with the class: water birds are usually depicted on top of a background with water. For evaluation, the classifier shown image of "water bird" on top of land background (and vice versa). Sim- ilar the RL setting, SL evaluation is done using pairs of inputs that are rarely seen together during training. How- ever, whereas the SL. setting aims learn classifier that ignores certain aspects of the input, the RL setting is dif- ferent because the aim is to learn a policy that can reason about both inputs.
1
6
Figure 9: SPURIOUS CORRELATIONS: To understand how spurious correlations are re- lated to stitching, let's look at this simple determinisite MDP in which three data col- lecting policies (red, blue and green) collect the offline dataset. Note that an OCBC al- gorithm which ignores the state, can also achieve zero training loss Eq. (6) this of- fline dataset. Whenever 4 the desired goal the dataset, action right is always optimal irrespective of the current state. Any SL al- gorithm that Iearns minimal dccision rule will in fact learn ignore the
There is another connection between the RL setting and spurious correlations, connection that makes the RL setting look the opposite of the SL setting. For some goal- conditioned RL datasets, the current state is sufficient for predicting which action leads to the goal - the policy does need look the goal. In other datasets, the goal is sufficient for predicting the correct action. However, for navigating between pairs of states unseen together in the dataset, a policy must look at both the state and the goal inputs.
reduce the training loss to zero this case (60) during test time, starting at state 2 and condi- tioned on goal 4 such an SL algorithm will ignore the current state and move towards right which clearly suboptimal.
B DIDACTIC EXPERIMENTS.
3
3
3
4
4
1
n
1 2 chonce
1
1
-
4
0 50%
4
0
2
5
6
2
6
1
1
chasc
c0%schance
1
v1
stochastic-v1
v2
Stochsstic-12
Figure 10: MDP with 5 states and 2 actions (up and right). All episodes end after taking two actions. Data collected using two policies (red and blue). only difference between v1 and vl-stochastic, is the data collecting policies are stochastic states 0 and 1. The stitching task navigate from states 1
Figure 11: MDP with 7 states and 2 actions (up and right). All episodes end after taking three actions. Data collected using two policies
and blue).
The only difference hetween v2 and v2-stochastic, is the data collecting policies are stochastic at states 0 and 1. The stitching task is navigate from states to 4.
Can OCBC algorithms exhibit combinatorial generalisation? In Fig. 5 we can see that OCBC algorithms do not always achieve zero performance. To understand why OCBC algorithms seldom two offline datasets collected from simple didactic MDP
perform some amount of stitching, (see Fig. I0). In v1. 3/10 random seeds of DT are successfully able to navigate from states 1 to while 7/10 fail. We show that in the both cases, the DT model picks up of two spurious
13

---

relations present in the dataset. In the success seeds, the DT model learns output action up and ignore the goal, whenever it is in state 1. In the failure seeds, the DT model learns to output action right, and ignore the state, whenever the goal is 4. In stochastic-v1, we deliberately remove the spurious relation that leads to success - the model can no longer ignore the goal, as the optimal action depends it, i.e., if the goal is then the optimal action is right, and if the goal is the optimal action is up. In stochastic-v1, we see that all 10 seeds of DT fail. To be certain that models ignore either states goals, we check that the model's outputs remain invariant them. This affirms the hypothesis that OCBC algorithms can succeed if the spurious relations in the offline datasets set them up for success.
Similar to the above experiment, we perform more experiments on two offline datasets v2 and stochastic-v2 (see Fig. 11 v2, 10/10 random seeds of DT are able to successfully navigate from states 0 to 4. The DT model in all cases picks up on the only spurious relation present in v2 the optimal action depends only on the current state and the goal can be ignored. We ensure that the model actually ignores the goal at state 0, by checking that its outputs remain invariant changing goals. In stochastic-v2, we remove this spurious relation; the optimal action at state depends on the goal as well, i..,, if the goal is 0, then the optimal action is up, and if the goal is 6, the optimal action is right. After removing the success spurious relation, we observe that the success of DT drops to 0/10 seeds. This result also algins with the same hypothesis that OCBC algorithms can succeed if the spurious relations in the offiine datasets set them up for success at test time.
C EXPERIMENTAL DETAILS
C.I ENVIRONMENTS
Goal conditioned environments. We use the "point" and "ant" mazes (umaze, medium and large) from D4RL [13. As discussed Section we carefully collect our new offline datasets to test for stitching combinatorial (see Fig. for visualisation). In the "point" maze, the task is to navigate ball with degrees of freedom that is force-actuated in the cartesian directions x and y. In the "ant" maze task, the agent is 3-d ant from Farama Foundation (62). To collect data for the "point" maze, we use a PID controller. To collect data for the "ant" maze, we use the same pre-trained policy from D4RL 15
In Fig. 12] provide visualisation of the offline dataset in all "ant" mazes.
Return conditioned environments. For the return conditioned tasks, create new task using Miniworld (58] (See Fig. 4). The task is to collect both the keys and return to the start position. A reward of is received after collecting each key. This task is image based and partially observable.
The agent recieves first person view of the world. At any time, it can choose amongst 5 actions [forward, backward, turn right, turn left, pickup). There are two data collecting policies, each collecting only one of the key. These data collection policies are implemented by controlling the agent manually.
C.2 IMPLEMENTATION DETAILS
In this section we provide all the implementation details as well as hyper-paramett parameters used for all the algorithms in our experiments DT, RyS, and RvS + temporal data augmentation.
DT. We used the exact same hyper-parameters that the original DT paper 2 used for their mujoco experiments. The original DT paper 2 conditioned the transformer on future returns rather than future goals. For our experiments, we switch this conditioning to goals instead. At every time-step the transformer takes in input the previous action, current state, and desired goal. The desired goal remains constant throughout the episode, but is still fed to the transformer at every timestep. All hyperparameters used for DT are mentioned in Table T
RvS. RvS is OCBC algorithm which uses fully connected neural network policy. We use the hyperparameters as prescribed by the original paper [3] All hyperparameters used for RyS are mentioned in Table
14

---

RvS temporal data augmentation. As mentioned in Algorithm given method to cluster states together, only requires lines of code to add the temporal data augmentation on top an OCBC method. We use the k-means algorithm from scikit-learn 663 with the default parameters to group states together. Adding data augmentation on top of RvS introduces 2 extra hyperparameters, which mention in 3 We do tune both of these hyperparameters in paper. Nevertheless, do ablate the choice of K in k-means.
Table 1: Hyperparameters for DT.
hyperparameter
value
training steps batch size
2 10
256
context len optimizer
5
Adam
learning rate
1 10 3
warmup steps
5000
weight decay
1 10-4
dropout hidden layers (self attention layers)
0.1
3
embedding dimension
128
number of attention heads
1
Table 2:
Hyperparameters for RvS.
hyperparameter
value
training steps batch size optimizer learning rate hidden layers
1 10
256
Adam
10-a
2
hidden layer dimension
1024
Table Hyperparameters for temporal data augmentation.
hyperparameter
value
K (number of clusters for k-means):
umaze
20
medium
40
large
80
(probability of augmenting goal)
0.5
C.3 DIFFERENCES BETWEEN ORIGINAL D4RL AND OUR DATASETS.
We made two main changes in the way our datasets (Fig. Fig. were collected compared the original D4RL datasets (Fig. 13). First, we ensure that different data collecting policies have distinct navigation regions, with only small overlapping region. This change helps to clearly distinguish between algorithms that can and cannot perform combinatorial generalisation. Second, the agent in the original D4RL datasets often moves direction that largely dependent on its current location in the maze. For example in the topmost row of the umaze, the D4RL policy always moves towards the right. To reduce such spurious relations, we randomize the start-state and goal sampling, for the data collecting policies. That is, in the topmost moves both towards the right or left, depending on its start-state and goal. Details about how such spurious relations can hamper combinatorial generalisation discussed in Appendix A
of umaze datasets, the data collecting policy
15

---

E
Umaze
Medium
Large
Figure 12: Offline datasets that we collect for the "ant" mazes. Different colors represent the navigation regions of different data collecting policies. See Fig. 3 for similar visualisation of the "point" maze datasets.
Umaze
Medium
Large
Figure 13: Similar Fig. and Fig. 12
visualisation of the original d4rl datasct. This not an aciual trajectories that are present inside those datasets,
exact visualisation collecting policies that those datasets used [[3]. During data collection, the policy starts from one starting region, which is marked by white star. The policy navigates to goal selected from one of the goal regions, which marked white circles. During test time, start-states and goals are selected from similar starting and goal regions, making difficult to evaluate the combinatorial generalisation offline algorithms.
visualisation the data
D PROOFS
D.1 PROOF OF LEMMA 4.1
We prove this Lemma by providing a simple counterexample. Consider the deterministic MDP shown in Figure which has five states [1,5] and two actions (right, up). The states four and five are absorbing states; once the agent enters one of these states, it will stay there for eternity. There are two data collecting agents 3 and 3h-2, which navigate upward from state two to state five, and rightward from state one state four respectively. Both policies collect equal amount of data p(h = = p(h = 2) = 0.5. The BC policy 9 (equation 5 is shown in purple to indicate that is obtained from combining data from both blue and red policies. We will prove that the training and testing distribution (LHS and RHS of Eq. 9 are not equal for the counterexample state-goo goal pair
{s: 2, 8t+ 4).
16

---

By(h) 1 |+ (sL+ =4 4| 8t = 2)pp 8m (st =2)
P4 (sL+ = 4 =4|s = 2) (st = 2)
= p+ (4 12)pt 22ppt=11 (2)
P+
Ah=2
(s++
8t = 2) 2
(s+ = =4 =4|s= == 2)(1 - ?)
2 (4
+
2
p+
1 (412)(1 -7)
Mn-
0
(1-2)² 7)² oc
Σ
P.
(4
+
2
t=0
(1 -)² 00
(1-))²
2 + 33
p'p -n-1 | (4
+ ....
2
2
2
t=0
(1 - 7)
(1 7)
x 0
0
2
The LHS and RHS are unequal for all values of E (0,
D.2 TEMPORAL. DATA AUGMENTATION DOES APPROXIMATE ONE-STEP STITCHING.
The reason why OCBC sampling does not lead to stitching is that it samples context ~ pp(p 8) at the start state, and commits to it by following 8n for the entire trajectory. On the other hand, the BC policy follows an average of all behaviour policies at every step of trajectory Cx (a (h) s)). One could think of the BC policy, as sampling a new context /t ~ p(h) at every timestep of the trajectory, and following the corresponding 8n. for that timestep. Intuitively, this ensures that all possible combinations of behaviors are sampled by the BC policy. These two ways of sampling trajectories: (I) sampling a context only once the start and committing to it or (2) sampling new context at every timestep, are two extremes, and one can think of intermediate ways of sampling trajectories. For example, sample a context h ~ p(h) and a time duration t ~ p(t), and follow 3n for timesteps "t", and then sample a new context ~ p(h) and follow 3Ú for the rest of the trajectory. Since this form of sampling trajectories changes context only once per trajectory, we call the resulting non- son-stttionayy policy as one-step stitching policy. If we do this process twice per trajectory, the resulting non-stationary policy is called two-step stitching policy. According to this notation, the BC policy can be thought of as per-step stitching policyFi We will prove that for any state-action pair, the support of the goals reached by n-step stitching policies as n increases, is non decreasing. Hence, sampling from the distribution of n-step stitching policies with a larger n, can sample unseen combinations of state-goal pairs. Finally, we will prove that applying data- augmentation once, under smoothness assumptions, is approximating the distribution of the one-step stitching policy.
Intuition.
N-step stitching policy. Before moving to the proofs, we need to define the distribution of future states sampled from the n-step stitching policy. Let p(t) be the duration sampling distribution which samples the timestep at which to change the context. We assume that p(t) has non zero support for all time-steps (f0, 1.... }). In our case, p(t) is the geometric distribution (geom(1 7)). Hence, for a state-action pair s, the distribution over goals reached by the n-step stitching policy is
po-s (g a.
/
Elhhi s, a)p Bmg (wi s, a
Ep(hs wi-1)p+ 9o: (w; w )
Wi.n An
hr
Here Wi is the sub-goal at which the th switch is made and W,+1 is defined as g.
Lemma D.I. For n E N. for all 8, a, the support of p (g 9, a) is at least equal to the support of prstee (g 8, a), assuming that the duration sampling distribution p(t) has non zero support for all timesteps (for example, geom(1 - -))).
'hh BC policy is stationary policy, as it samples new contexts every timestep. There is no dependency on time.
17

---

Proof. Assume that goal-state g belongs to Supp sep (g a). We will prove that this implies [g s, a) we can assume without loss of
that 9 E Supp p0rl- s, a). Because g E Supp p p generality that 9 is visited l time-steps after switching the context n times.
Before switching for n+ 1shh time, the distribution of states visited by the n 1-step policy is identical
to the distribution of n-step policy. The + 1 switching duration is drawn from fn+1 p(t). Because p(t) has support over all timesteps ((t) = {0 + NJ), we know that p(tn+1 > k) > 0. That is, the probability that n++I-step policy does not make the n+ 1th switch for atleast k timesteps is non zero. This proves that for any finite value of k. there is a non zero probability that the distribution of states visited by the n + 1-step policy is identical to the n-step policy for atleast & timesteps after the nthh switch. Hence the goal state can also be visited by n + 1-step policy. Hence proved that Supp ( (g] 8, s,a) C Supp p0 88 a).
Using the above lemma, we can say that Supp Pd se (g 8, a) Supp Suppp"*l-ang (tllss 8, a) C Suppp s,a).... ...
(10)
Hence proved that for any state-action pair, the support of the goals reached by n-step stitching policies as n increases, is non decreasing. And from 4.1, we know that there can be states which are only visited by {n > St-s stitching policies. Hence, it is possible that the relations in Eq. (10) are strict.
Temporal data augmentation approximates one-step stitching policy. Under smoothness as-
sumptions for the discounted state occupancy distribution of the data collecting policies, we can show that temporal data augmentation Fig. approximates the distribution of one-step stitching policy. Specifically, we assume that for all s, a, g pairs and all data collecting policies Br. p (g 1 is L Lipschitz
8r (g s.0) p (w s.a)) S L(llg- - (w(l)
(11)
This is an important assumption for temporal data augmentation it ensures that states which are close together have similar reachability. Prior work has also studied this assumption [64] and applied it to practical settings (67]. Finally, our theoretical analysis uses a form of temporal
data augmentation which clusters states only with their nearest neighbour. That is, each group in
Algorithm D contains atmost unique states.
Lemma D.2. Under the smoothness assumptions above, all 8, pairs, temporal data augmentation p jemp-aus (g 8, a) approximately samples goal according the distribution of one-slep stitching policy
(t (g s.))).
Proof.
pemp-ma8 s, a)
Σ p(h| s, ap temp
u "w s,))) Σ p(h) w)p (g w)dw
6
p(h | ww w1) (g) wi )dw! Ep(h1s,a)lef* (wi 8, EL) Σ
an
h
Σ (wi s, Σ
p(h
wr 1a (g) wiwdwi = eL f
Ep(h p(h wi)r I w2 )dwi
m1 h
h
h
Ep(h s, (wi |s a) Σ
p(h
we w1)p 3/ (g wj)dii = O
101
h
isteppp (g
Here é is the maximum cutoff distance used by the temporal data augmentation to group states together.
18

---

(a) is followed by the fact that temporal data augmentation first sample waypoint w. and then sample a future goal from the trajectory containing w. This new trajectory can correspond to any data collecting policy Bn with probability p(h w). (b) wi is the initial goal sampled by temporal data augmentation, which is then substituted by its nearest neighbour w. It follows from Eq. (IT) that 8i
8h (wi s, a) S p (w s, a) cL, as epsilon the maximum possible L2 p+ (w s, a) eL S p+
norm between W1 and w.
19


Reviews:
[{"review_id": "vRvjgb56WZ", "reviewer": "Reviewer_dPJE", "summary": "First, the paper identifies one important capability present in dynamic programming but lacking in supervised learning approaches to RL: the stitching property, which generalizes between different trajectories. - The second conceptual leap connects the stitching property to a certain kind of generalization. Finally, the authors propose to outfit data augmentation from supervised learning for reinforcement learning as a simple method to improve generalization. Experiments on two enviornments show that RvS with state and goal augmentation improve over RvS and DT. It is also claimed that more data did not help the DT improve, nor did increasing the number of layers.", "questions": "- 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 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?\n- 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.\n- 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 generaliation would be equivalent to iid generalzation. But the problem is, in a sense, limited data: the context is never accessible and hence the data (or information) is fundamentally limited.\n- 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 fidning a better policy necessitate stitching.\n- 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.\n- 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?\n- 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.\n- Section 6 (More data): This seems at odds with condition 1 under \"popular datasets do not evalute stitching\". Surely, if every state and goal pair were included in the dataset, then DT would improve?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The theoretical foundation developed in this paper is interesting. There is high value in formalizing stitching, and hypothesizing on fundamental limits of certain algorithmic approaches that apply supervised learning to reinforcement learning problems. I also like the ethos of the approach, bringing data augmentation to RL for a specific purpose", "weaknesses": "- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PBJebqXKe7", "reviewer": "Reviewer_DGi6", "summary": "The paper explores the differences between dynamic programming-based reinforcement learning (RL) and supervised learning (SL)-based RL. A key property of the former is \"stitching\" past experiences to address new tasks. The study relates this ability to a unique form of generalization. Through experiments, the authors reveal that SL-based RL methods might lack this stitching capability. To address this, they introduce temporal data augmentation, enabling SL-based RL to handle unseen (state, goal) pairs effectively.", "questions": "Q1: 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?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper provides a comprehensive analysis of the differences between dynamic programming-based RL and SL-based RL through the lens of generalization, offering insights into their inherent properties.\n2. A novel augmentation method is proposed to make SL more generalizable.", "weaknesses": "1. 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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lu2YAfFZwa", "reviewer": "Reviewer_4eRX", "summary": "This work analyzes SL-based RL approaches for (goal-conditioned) offline RL, where data are collected using various policies and one hopes to \"stitch\" existing experience/trajectories to generalize to unseen (start, goal) pairs. It shows that common outcome conditional behavioral cloning (OCBC) methods can fail this task and shows a lemma (Lemma 4.1) that learning from training experience with different contexts may not lead to generalization (in terms of the mixed behavior policy). Then the paper proposes to use temporal augmentation that essentially stitches two trajectories together to have better coverage for the (start, goal) pair. Experiments show that the augmentation can be effective for some offline RL tasks.", "questions": "Q1: Any suggestions for a reasonable distance metric for the states? \n\nQ2: Are there any theoretical justification for the temporal augmentation?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Studies an important problem of generalization from offline data\n- Shows a counterexample of the training and test discrepancy (Lemma 4.1)\n- Clear writing in most places", "weaknesses": "1. Some key concepts and results are not clearly explained.\n\n    \n\n    1.1. 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.\n\n    1.2. \"Our paper focuses on the problem of stitching generalization, not finding optimal policies; however, stitching generalization should be a necessary component for finding optimal policies.\" Proof needed.\n\n2. 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.\n\n3. 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). \n\nMinor\n- $p_+^{\\beta_h}(s,a)$ in Eq.(7) is undefined.\n- Both generalization and generalisation are used.\n- point-maze large taks -> task", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ToCH09j43m", "reviewer": "Reviewer_6ajj", "summary": "* This work discusses a form of generalization in reinforcement learning (RL) called \"stitching generalization.\" According to the paper, the stitching generalization cannot be achieved with existing supervised-learning (SL) based RL algorithms because it is fundamentally different from the generalization on independent & identically distributed (i.i.d.) dataset. This is explained via a concrete example of an MDP depicted in Figure 1.\n* Also, this work propose \"temporal data augmentation\" technique for goal-conditioned RL to actually implement stitching generalization inside the SL-based RL algorithms. This technique is based on careful sampling of augmented goal, which is determined by *closeness\" of states based on a clustering algorithm.\n* Moreover, this work provides a novel benchmark to evaluate the stitching capability of RL algorithm.\n* Lastly, this work empirically show the efficacy of the temporal data augmentation + outcome conditional behavior cloning (OCBC) algorithms on their own benchmark.", "questions": "* 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)?\n* In page 5, what does it mean by ‘combinatorial generalization’? Do you have a definition or any reference for that?\n* 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)$ ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The motivation is clear and the main message is well presented. I enjoyed reading the paper.\n* Their main method of data augmentation is novel and intuitive.\n* The empirical results are promising, even after making the task more difficult to generalize.\n* Lastly, the limitations of their work is clearly mentioned. I also look forward to scalable version of their algorithm.", "weaknesses": "* **Theoretical contribution is a bit marginal.**\n  - 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.\n* **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?**\n  - 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.\n* **Weakness on Algorithm 1 (OCBC + Temporal Data Augmentation) regarding additional hyperparameter**\n  - 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.\n* **It seems necessary to conduct additional Experiments on DT + temporal augmentation.**\n  - 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.\n* **There seems to be some minor typos.**\n  - pg 1. Kaelbling (1993) $\\rightarrow$ (Kaelbling, 1993)\n  - 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)]$\n  - pg 3. “… and then sampling a trajectory from the corresponding policy $\\\\{\\beta(a\\mid s, h)\\\\}$.” $\\rightarrow$ $\\beta(a\\mid s, h)$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)