Return Augmentation gives Supervised RL Temporal Compositionality
Authors: Keiran Paster
Silviu Pitis
Sheila A. McIlraith
Jimmy Ba
Venue: Submitted to ICLR 2023 (2023)
Keywords: reinforcement learning
offline reinforcement learning
decision transformer
behavioral cloning
dynamic programming
data augmentation

Abstract: Offline Reinforcement Learning (RL) methods that use supervised learning or sequence modeling (e.g., Decision Transformer) work by training a return-conditioned policy. A fundamental limitation of these approaches, as compared to value-based methods, is that they have trouble generalizing to behaviors that have a higher return than what was seen at training. Value-based offline-RL algorithms like CQL use bootstrapping to combine training data from multiple trajectories to learn strong behaviors from sub-optimal data. We set out to endow RL via Supervised Learning (RvS) methods with this form of temporal compositionality. To do this, we introduce SuperB, a dynamic programming algorithm for data augmentation that augments the returns in the offline dataset by combining rewards from intersecting trajectories. We show theoretically that SuperB can improve sample complexity and enable RvS to find optimal policies in cases where it previously fell behind the performance of value-based methods. Empirically, we find that SuperB improves the performance of RvS in several offline RL environments, surpassing the prior state-of-the-art RvS agents in AntMaze by orders of magnitude and offering performance competitive with value-based algorithms on the D4RL-gym tasks.

RETURN AUGMENTATION GIVES SUPERVISED RL TEMPORAL COMPOSITIONALITY
Anonymous authors
Paper under double-blind review
ABSTRACT
Offline Reinforcement Learning (RL) methods that use supervised learning or se- quence modeling (e.g.. Chen et al. (2021a)) work by training return-conditioned policy. A fundamental limitation of these approaches, as compared value-based methods, is that they have trouble generalizing to behaviors that have a higher return than what was seen training (Emmons et al., (02T). Value-based offline- RI. RI. algorithms like CQL use bootstrapping to combine training data from multiple trajectories to learn strong behaviors from sub-optimal data. We set out to endow RL via Supervised Learning (RvS) methods with this form of temporal composi- tionality. To do this, we introduce SUPERB, a dynamic programming algorithm for data augmentation that augments the returns in the offline dataset by combin- ing rewards from intersecting trajectories. We show theoretically that SUPERB can improve sample complexity and enable RvS to find optimal policies in cases where it previously fell behind the performance of value-based methods. Empir- ically, we find that SUPERB improves the performance of RyS in several offline RL environments, surpassing the prior state-of- of-the-art RvS agents in AntMaze by orders of magnitude and offering performance competitive with value-based algo- rithms on the D4RL.-gym tasks (Fu et al.) 2020).
INTRODUCTION
The use of prior experiences to inform decision making critical to our human ability to quickly adapt to new tasks. To build intelligent agents that match these capabilities, it is natural to seek algorithms that learn to act from preexisting datasets of experience. Research on this problem, for- mally known as offline reinforcement learning (RL), has focused two main approaches. The first takes existing off-policy RL algorithms, such as those based on Q-learning, and alters them to reduce issues caused by distributional shift. The resulting algorithms use value pessimism and policy constraints to keep actions within the support of the data distribution while simultancously optimizing for high returns (Fujimoto et al. 2019 Kumar al. 2020a). The second, RL via Super- vised Learning (RvS), draws inspiration from generative modeling and supervised learning to learn eercome-conditioneee policy models and uses them predict which actions should be taken in order to get a high return (Schmidhuber) 2019 Chen et al., 2021b Emmons et al. 2021).
RvS algorithms are appealing due to their simple training objective, robustness to hyperparame- ters, and strong performance, especially when trained in multi-task setting. Recently, however, attention has been brought to their suboptimality in certain settings, such as stochastic environments (Paster et al. 2022, Villatlor et al., 2022] Eysenbach et al.) 2022) and, as we remark in this work, offline settings where temporal compositionality" is required for good performance (see, e.g., Fig- ure Tp. Because dataset, they may fail to extrapolate to higher returns not present in any single trajectory but made possible by combining behaviors from multiple trajectories. While value-hased methods use dy- namic programming to compose behaviors across trajectories, RyS approaches fail to fully exploit the temporal structure inherent to sequential decision making (Sutton 1988).
couur-conditioned RvS agents are trained using returns calculated from offline
In our work, we aim to endow RyS methods with this kind of temporal compositionality. Tradition- ally, RvS agents are trained using return labels computed by summing over the future rewards of
We define temporal compositionality the composition of hehaviors from different timesteps within trajectory from different trajectories.

---

S
r=
F= 10
10
10
2
15
r= 16
r= 10
r=2
r= 10
10
2
0
r=
r=0
r=
r=
a:
So
Motivating Example
Original
Augmented Return-to-gos
Return-to-gos
this illustrative offline RL task, the dataset consists trajectories gathered by two suboptimal policies: the bluc/solid policy, which goes to the left and either gets return of 12 or 15, and the red/dashed policy, which goes to the right and gets a return 12. Middle: Typical RvS algorithms compute the return-to-go by summing up the rewards along each empirical trajectory. While the optimal trajectory (Su. $1, 83, 84) has a value of 20, no empirical trajectory has value of 20. Right: Our method, SUPERB, combines rewards from multiple trajectories when calculating return-to-go labels and applies a label of 20 to each state and action along the optimal trajectory.
Figure Left:
each trajectory individually. Our main insight is that we can apply the n-step temporal difference (TD) relation (Sutton Barto 2018) to offline data augment the return labels used in training with returns that are made possible by composing different trajectories. While Q-learning algorithms aim to learn the value of optimal policy, our method uses distributional value function to assist in generating distribution of feasible return labels which agent. We propose to iteratively apply our method to augment returns, which we show can generate exponentially more labels for Rvš agents train Our main contributions
train return-con coun--onddtiooed RvS
We show that there are some environments where, both analytically and empirically, return-condin conditonnd RvS agent falls behind value-based agents that use bootstrapping, and propose a data augmentation method that bridges this gap by using the n-step TD relation combine rewards from different trajectories.
2. We show that temporal compositionality can exponentially improve the number of training trajectories and is necessary to learn optimal policies in some offline RI. datasets.
3. We evaluate our method the D4RL offline RL suite. Our method dramatically improves the performance of RvS on AntMaze environments, where optimal policies must stitch together behaviors from partial demonstrations, from near zero to state-of-the -frtt. We also obtain performance competitive with value-based methods on the D4RL tasks.
METHOD
2.1 PRELIMINARIES
We consider the offline Reinforcement Learning (RL) setting where the environment is modeled as a Markov Decision Process (MDP), MMM A, T,, R, 7), consisting of state space, action space, transition function, reward function, and discount factor, respectively (Lange et al., 2012 Sutton & Barto 2018}. The agent is given fixed dataset D of state-action-rewarr trajectories (31, d1 Tı, 82, 02, T.....) produced by Enon-Markovina "empirical policy" ग. acting in M, and that obtains high return E. rr2 when executed in M.
tasked with learning policy ma
2

---

We build on the Reinforcement Learning via Supervised Learning (RvS) paradigm (Chen et al. 2021b Schmidhuber 2019). in which the agent learns a return- condittond policy Pa(a s, 2), where 2 an empincal statistic about trajectories that start in state (s, a), and then conditions this policy on a favorable =" to obtain its policy [a(a]] () po(a 8, ). We condition on the original discounted return objective = 9 ". 7"re, and leave extensions other potential 2 future work. We define a return function G T return. pa trained to model empirical samples (s, ()
- IR that maps from (sub)trajectories to their discounted D using maximum likelihood estima- tion, where (s, a) are sampled uniformly from all (s, a) pairs, and G((Ta) ) is the return of the empirical subtrajectory that begins with (s, a). As in prior works (Chen al. [201a] Emmons al., 2021 is tuned along with other hyperparameters to maximize performance.
2.2 BOOTSTRAPPING
Many of the most successful RL methods such as Q-Learning (Watkins Dayan 1992) exploit the time structure of MDPs to do temporal difference learning (Sutton) 1988 Tesauro 1995), form of dynamic programming where a learned model of the expected return is used to provide learning targets, either recursively for itself or for other agent modules. This process, which has come to be known as "bootstrapping", allows the agent to learn off-policy and achieve higher returns than those observed by the empirical policy. One way understand bootstrapping is that it allows the agent to implicitly stitch together empirical experiences via the learned return model (Sutton) 1988). This effect particularly visible the offline RL setting, where Q-learning based methods such BCQ (Fujimoto et al. 2019) and CQL (Kumar et al. 2020a) are able to far surpass behavior cloning (BC) and obtain higher returns than the empirical policy TT..
Unlike Q-learning based methods, the ordinary approach to training Pe in RyS methods does not use bootstrapping: g is simply computed as the cumulative (discounted) reward of the trajectory starting in (s, a). For this reason, the performance of RvS methods is fundamentally limited by the return of the best performing trajectories in D, since these upper bound the set of returns {G(T:)) used for training Po. While extrapolation via function approximation is possible and conditioning on the highest observed returns allows RyS agents to outperform the average empirical policy, typical RyS training ignores the known temporal structure of M and lags behind bootstrapping methods on certain datasets (Emmons et al. 2021).
On other datasets, however, RyS methods still outperform Q-learning based methods and are known to be more stable when using varied hyperparameters (Emmons et al./ 2021] Ghosh et al.] 202T). One potential reason for this is that (1) the superviie bootstrapping used by Q-learning, when combined with (2) function approximation and (3) off- off-policy data, suffers from general source of training instability known as the "deadly triad" (Sutton & Barto 2018). By contrast, Rv$ methods use (by assumption) unbiased return labels generated by Monte Carlo rollouts of ne.
2.3 SUPERB: SUPERVISED RL WITH BOOTSTRAPPING
To resolve a fundamental limitation of RyS methods and grant them the ability surpass the highest returns in the empirical dataset, we propose to augment the empirical data D with "bootstrapped" returns calculated by combining multiple trajectories. Recall that in return return-condee conditioned RvS, training data consists of (s, a, g) triples, where g is possible return that can result from taking action a in state s. We augment the dataset with novel possible returns by observing that the return G(T 72) of two concatenated subtrajectories can be computed from parts using the n-step TD relation:
G(TT22 = G(n) + 7 -
To generate returns that could occur by stitching two subtrajectories together, we model the distri- bution of the returns of trajectories that start at the last state of TΙ. and create augmented returns by replacing G(T2) in Equation with samples from this distribution. Our method, which we call SUPERB, for Supervised RL with Bootstrapping, consists of the following phases:
Learning a Model of Returns. The input our algorithm is dataset of states, actions, and returns (8, 9) ~ D. As the first step towards augmenting the returns, we learn a model (G of the distribution of returns possible from any given state in the dataset, parameterized by ø.
Augmenting the Dataset with New Returns. In the second phase of our method, we strive to create novel return labels that represent possible values agent could achieve from each (s, a) pair
3

---

Algorithm 1 SUPERB
function SUPERB(data D {r), steps N)
function AUGMENTDATA(D {r}, model Vo)
2 D
D
0
7 1. ............... do
for m ....M do
V < TRAINREWAKUMODELLD"""
Sample ~ D Choose prefix T1 = as, of) ().
D AUGMENTDATA((" v)=-1)
g = G(T)) + + 9 ~~ Volo')
Pa TRAINRYSAGENT(D") return Po
D' + (,,00''
return COMBINE(D, D')
function TRAINREWARRDODELL(D {s, {a.g))
Sample (8,9) ~D
function TRAINRYS = ((s, = a, a,
Train (G(G using MLE
Sample (s, a, ~ D
return VP
Train pa 8, using MLE
return yos
in the dataset. To do this, SUPERB chooses a prefix for each trajectory in the dataset, denoted Fi (s, a, T..... s'), and combines with another subtrajectory in the dataset imagine a hypothetical new trajectory. Rather than explicitly combine trajectories and calculate returns, SUPERB uses the learned return model. Specifically, rather than calculate return associated with (s.a) D as G(Txa) E,7/ra as in prior couut-conddiioned RvS works (Chen et al. 2021a; Emmons et al., 2021]. SUPERB uses the n-step TD relation Eqqution I] to calculate the return of an imagined, combined trajectory as G(TTF)
G( + > go', where Is', sampled from the learned return model (G s'), represents the value of possible subtrajectory that starts in s'. These returns are added to the new, augmented dataset D'.
Performing these two operations only once gives us returns that are stitched from at most two sub- trajectories. To generate labels that compose 2 subtrajectories, we repeat this process N times.
Training the Return- Conditioned Policy. Finally, return couu--onditioned RvS policy pa s. s.9) is learned via maximum likelihood from D' as done in prior works (Chen et al. 2021a Emmons et al, Kumar et al. 2019). The full process is summarized in Algorithm |
2021
Schmidhuber 2019
2.4
BIASING SUPERB TOWARDS RETURN MAXIMIZATION
The implementation of the different phases in SUPERB involves several important design choices: what distributional return model to use for (G s), what distribution over possible T₁ to use when forming the bootstrapped returns, how to sample from (G s), and how much of the original dataset D to keep. We opt to inform our choices based on the downstream use of the agent: return maximization. In particular, since the RvS policy typically conditions on high returns, we make several choices to maximize the performance of the agent in this regime.
First, rather than sample directly from (G s) when augmenting returns, which just as likely to form low returns it is high returns, we learn (o 8) using quantile regression in Dabney al. (018)) and bias our samples by only drawing from the upper quantiles (specifically, we use the mean of the top 5 of 20 quantiles). Second, when choosing 7i, which represents the first part of trajectory to be spliced together with a return from a different trajectory, we use a backward induction approach consider all possible T1. only choosing augment using 71 when the augmented return is higher than the existing return. Finally, since as a consequence of this decision the augmented returns are guaranteed to be higher than the original returns, we choose to discard the returns in the original dataset when training the RvS policy. See Appendix B for further details.
These design choices introduce some similarity to Q-learning methods. However, while Q-learning aims to find Q-function for the optimal policy by minimizing the Bellman error, the objective of SUPERB is to alter and expand the distribution of returns on which an RvS agent is trained. Despite these different motivations, biasing SUPERB toward return maximization allows it to be understood as way of unifying RvS and Q-learning (cf. TD(A), Sutton (1888) ): as the number of bootstrap iterations increases, the augmented dataset produced by SUPERB becomes increasingly similar to sample-based representation of an optimal Q function. By using an intermediate number bootstrap iterations, we hypothesize that we can trade-off the strengths of RvS and Q-learning (i.e., training
4

---

$0
82
8NI1
Figure 2: N-Step Chain MDP: The MDP is initialized in so and terminates in 8 with actions always leading to the next state in the chain. this MDP there are JN possible trajectories.
stability in case of no bootstrapping vs. the ability stitch together large number of behaviors). Our intermediate approach may be helpful in many realistic offline RL datasets (such as those found in D4RL (Fu et al. 2020)), which only require a small amount of stitching to find strong behaviors.
2.5 ANALYSIS
In this section, we show formally how temporal composition is necessary generate optimal reward labels for some sets of offline RL problems, as well as how it can generally improve the sample complexity of conditioned RvS methods.
Proposition 1 (Necessity). For any positive integer n, there exists an MDP M and data generating policy Te for which n-step bootstrupping is strictly necessary to generate dataset 'D containing (s, ",,.gg)), where g* is the maximum reward achievable for a trajectory starting in (s,a).
Proof. See Figure the case 1. For the general case, one construct MDP and offline dataset where non-Markov Te is generated by N + different policies, each containing a segment of the optimal policy. The details are found after the proof of Proposition Za in Appendix A
Note that the proof of Proposition explicitly constructs non-Markovian me. Some thought reveals that non-Markovian Te is required for necessity: if we continually sample data from Markovian Tex we will eventually sample the highest possible return that could be reached by augmenting returns produced by we- This suggests that we should expect SUPERB to provide more benefit when augmenting datasets produced by Mar-ooaon policies Te, such as the case when the dataset produced by a mixture of policies.
The following Proposition (and Proposition 2A in Appendix A assumes an "N-Step Chain MDP", exemplified in Figure 2 with N + 1 states, and b actions in each state. The Chain MDP simplifies analysis in a way that is particularly favorable toward bootstrapping. Although the results will not hold generally, they demonstrate that bootstrapping has the potential be extremely valuable in environments where trajectories have many possible points of overlap. Indeed, in the Chain MDP, every trajectory overlaps with every other trajectory every state. On the other end of the spectrum, we could consider an MDP where the only point of overlap is in the initial state: in such an MDP, there would be no advantage bootstrapping.
Proposition 2 (Coverage of Trajectory Space). In deterministic N-step Chain MDP with branch- factor b, it is possible for N-step bootstrapping to (implicitly) capture full "coverage" of the size B trajectory space with only b empirical trajectories- exponential increase in coverage relative to no bootstrapping. Coverage here refers to the percentage of unique trajectories present in the dataset, where we consider two trajectories with the same actions taken in every state as equivalent.
Proof. We apply our Augmentation Bootstrapping
nugmeentaioo-Booosttapnnngg Equivalence Lemma, found in Appendix A
which states that return label bootstrapping is equivalent to augmenting the set of trajectories by stitching together subtrajectories. While the equivalence itself is interesting, a formal statement of the lemma is long and relegated to Appendix A
Our analysis suggests that SUPERB can offer nontrivial advantage over plain RyS methods; how- ever, the maximization bias we introduce via our design choices, together with the approximation error introduced by model of the return distribution, suggests tradeoff between the benefit of bootstrapping and risk of increasing bias. As a result, the right number of bootstrapping iterations may be problem dependent. In the next section, explore how SUPERB performs empirically.
5

---

3 EXPERIMENTS
We designed our experiments answer the following questions:
Can SUPERB improve RvS performance on environments requiring temporal composition?
How does performance vary with the number of bootstrapping iterations?
What design choices are important when doing return augmentation?
Key experimental details are described in the main text, with minor details left to Appendix B
3.1 BENCHMARK TASKS
We conduct our experiments on the D4RL benchmark suite (Fu et al. set by prior works in offline RL. In particular, we focus on the Gym MuJoCo tasks (Brockman et al., 2016), which include Hopper, HalfCheetah, and Walker2d, as well as the AntMaze tasks ("umaze", "medium", and "large" mazes). The gym tasks are designed to test how offline RL algorithms deal with suboptimal agents and narrow data distributions, while the AntMaze tasks test their ability to learn from sparse reward data generated by non-Markovian policies. We are primarily interested in performance on the AntMaze tasks, where strong performance requires stitching together behaviors from different trajectories, the vast majority of the datasets consists of trajectories that do not start at the initial state and travel to the goal. As reported by Emmons et al. (2021). RvS methods perform poorly on these tasks without additional information (e.g., RvS-G in Emmons et al. (2021) does goal relabeling using the knowledge of the state dimensions that correspond with the goal state).
2020), following precedent
3.2 BASELINES
We base our implementation of RvS on the code provided by Emmons et al. (0021) and rerun the RvS baseline with any improvements we made to the codebase to ensure a fair comparison. We compare with behavioral cloning and filtered behavioral cloning (Chen et al., 2021a), as these have similar implementations to RvS and provide a simple sanity check. We also compare with Decision Transformer (Chen et al., 2021a). statc-of-the-arn implementation of RvS which uses transformer (Vaswani et al. (217)) as a policy. Finally, we compare with several Q-learning variants. We compare our method with CQL (Kumar et al. (2020b), which is the standard for many offline RL tasks, and IQL (Kostrikov et al. 2021), which has similarities with our method discussed in Section 4 We report results for CQL, BC, %BC, DT, and RvS as reported in Emmons et al. [202I) For IQL (Kostrikov et al., 2021). we re-run their AntMaze experiments of the environments and use their originally reported results otherwise. To ensure fair comparison, rerun RvS using our hyperparameters and report the results under the name Rv$ (Ours)
the updated v2 version
3.3 PERFORMANCE ON D4RL TASKS
Table 4 shows a comparison of the performance of our method with prior methods. Our experi- ments reveal that SUPERB outperforms the base RvS agent across almost all tasks, especially on the medium and large AntMaze environments, where the augmented returns generated by SUPERB boost performance to be state-of-the-art. In contrast with prior works (e.g., [Emmons et al. 2021] Janner et al. (2021)), which use goal-state information to achieve high performance in AntMaze, SUPERB achieves strong performance only using basic information from the reward signal. On the D4RL gym tasks where temporal compositionality is less important, SUPERB generally retains the performance of Rv$ without any augmentation. Consistent with our hypothesis from (section 2.5) that SUPERB provides the most benefit datasets formed by non-Markovian ne. non-trivial im- provement is obtained on two out of the three the medi um sumuueee replay datasets, which are the only non-Markovian gym datasets where expert performance is not obtained.
Our results are more modest on the other D4RL-Gym tasks, where the datasets are single-task and, in many cases, only collected with a single Markov policy. We believe, however, that most real- world offline-Ri datasets are likely to be multi-task and gathered by a mixture of different policies, where SUPERB B will have large effect.

---

Task
CQL
IQL
BC
% BC
DT
RvS-R
RvS
SUPERB
antmaze-umaze-vv
44.8
83.7
54.6
60.0
65.6
64.4
84.2
85.4
23.4
66.7
45.6
46.5
51.2
70.1
77.7
76.9
antmaze-mediun-plan
play-v2
0.0
74.7
0.0
42.1
1.0
4.5
2.7
76.5
antrarz-meedium-ddieeee-
0.0
71.0
0.6
7.7
2.7
78.1
autmazz-large-plaayy22
0.0
44.3
28.0
0.0
3.5
0.8
25.8
antiaaz-large-divesee-22
0.0
49.0
0.0
34.3
0.2
3.7
6.9
36.9
hallcheetah-ramlom--22
18.6
2.3
2.0
2.2
3.9
2.2
2.2
hopper-tandom--2 wa walker2d-randoom-22
9.3
4.8
4.1
7.5
7.7
5.8
14_3
2.5
1.7
1.7
2.0
-0.2
6.0
5.6
hal lcheehhhhhhhee weweeeaaah---nueeeeeeee (uedium-replay-v
47.3
44.2
36.6
40.6
36.6
38.0
39.1
39.1
hopper-mediune popper-medium--eplla--2-
97.8
94.7
18.1
75.9
82.7
73.5
77.9
90.3
walker2d meld---mdi----mm repaayyy--pp22p2
86.1
73.9
26.0
62.5
66.6
60.6
50.5
halfcheetah-mmn meneruun--mmiim-
49.1
47.4
42.6
42.5
42.6
41.6
42.3
42.1
Kopper-medium-vV walker22dmmechum-222
66.3
52.9
56.9
67.6
58.9
56.8
82.9
78.3
75.3
75.0
74.0
71.7
72.4
74.7
saffehhetaa-mmddim--xpppep
85.8
86.7
55.2
86.8
92.5
hepper-medium-eepeere
102.0
91.5
52.5
110.9
107.6
101.7
109.8
110.1
wulkerr2U-mediummexperttt-..
109.5
109.6
107.5
109.0
108.1
106.0
106.9
107.9
Table I: We evaluated our method with two bootstrap iterations on the AntMaze and Gym tasks from D4RL (Fu et al. 2020 Across the AntMaze tasks, which require stitching together behaviors from several suboptimal trajectories in order to perform well, SUPERB provides strong performance improvement over the baseline RvS agents that don't use bootstrapping. outperforming the prior state-of-the-an in 3 of the 6 mazes. The benefits of SUPERB the Gym tasks more modest, though we observe improvements over RvS on a few the suboptimal datasets, particularly those where the dataset
generated by non-Markovian combination of policies.
best
method, underline denokes the best RvS method. Results averaged seeds.
Betirn-oggu
Auguettal: Tur
80
S
S1
70-
Normalzed Doffrmane
SteperiR(3)
60-
Augrented: Tter
Auguenenl:: Tue
wellac of
(b)
Figure Left: In the AntMaze tasks that include data generated by suboptimal agents, Su- PERB provides strong benefit. After two iterations of SUPERB, performance starts to decrease, indicating that the benefit provided by the extra bootstrapping step is outweighed by the ad- ditional errors added from function approximation. Right: We plot the augmented returns in antmaze weneeee-eiiiie un-diverse
for up to three iterations in order to visualize the effect of Su- PERB the maximum values associated with the trajectories in the dataset. Each pixel is colored according to the highest augmented return across observed states that fall within that pixel, default- ing Vnin if no states were observed (possible returns are in [Vmin: Vinaz] [-1/(1 0.997), 0)) Even single iteration serves to make the coverage of successful examples significantly more robust.
3.4 PERFORMANCE VS. BOOTSTRAP ITERATIONS
SUPERB can be applied iteratively N times to create return labels that are bootstrapped N times, equivalent to combining returns from up to + different trajectories. We ran SUPERB for up to three iterations on AntMaze in order to study the effect of this hyperparameter. Figure 3a/shows that
7

---

Ablation Study
SB(1)
SB(2)
SUPERB (Ours)
55.8
78.1
Remove Reward Transformation & Discounting
7.1
3.8
Remove Dynamic Target Return
41.67
32.1
Keep Original Returns from D
23.7
22.1
Table 2: In our ablation study -lt of modified reward function (matching the implementation in the Q-learning based baselines), dy- namic target returns, and reducing the dataset size by pruning the original return labels is necessary to achieve strong performance. Note that several of these ablations are still far stronger than the other strongest RvS method on this task, RvS-R, which only gets score of 7.7.
we find that combination
tasks that benefit from bootstrapping (the medium and large AntMaze environments), most of the performance benefit is present after only 1-2 iterations. Figure 3b shows the maximum return label for each state in the antmaze-med di um-divieee wwwidiivvveeeeee task. Our visualization of the returns used to train prior RvS methods (such as RvS-R. DI, and SB(O) in Table 4) in |Figurr 3b (top left) shows that the dataset contains examples of successful trajectories, but that these trajectories are sparse. Just or two iterations of SUPERB makes the coverage of successful trajectories far more robust, and we suspect further iterations only serve to add noise to the dataset. This empirical finding is reassuring since it suggests that in practice even one iteration of bootstrapping with SUPERB can provide a large benefit in practical datasets. However, as discussed in Section 2 one can imagine tasks more similar to the chain MDP which would benefit from many more iterations, meaning that there is no universally optimal number of iterations.
3.5 IMPORTANT DESIGN CHOICES
As discussed Section 2.4] there are several important choices when implementing SUPERB. Ta-
ble 2 shows the result of our ablation study the wwmmmazeeedddne verse-v2 environment where we empirically verify that each of the following elements is important.
Reward Transformation and Discounting in AntMaze. In the AntMaze environment, it is com- mon for value-based offline RL. algorithms (such as IQL. (Kostrikov et al.] 2021)) to apply a trans- formation to the reward to convert it from sparse reward of 1 when the goal is reached to reward of O when a goal is reached and -1 each step otherwise. We found that applying this reward trans- formation combined with high discounting factor of 0.997 improved performance of SUPERB, but did not improve the performance of our RvS baseline.
Dynamic Target Return Selection. RvS policies need to be conditioned on target return. In most prior work, this target is manually chosen and updated throughout the episode to account for rewards the agent receives. We take inspiration from Lee et al. (2022) and condition on the highest return quantile of our learned model (G s) for the current state during the course of an episode. We found this improves performance in some tasks, as compared to manually choosing the target.
Discarding Original Returns. In traditional data augmentation methods, the original data is kept as a part of the dataset. However, we found discarding all but the most bootstrapped returns can improve performance dramatically, perhaps due to easier policy learning.
4 RELATED WORK
Offline Reinforcement Learning In the Offline RL setting, an agent must learn policy that maximizes returns from static dataset of trajectories (Lange et al., 2012). When doing RL with offline datasets, the goal is typically to compose behaviors seen in the offline dataset while avoiding taking out-o--distribution actions that may lead to unexpected outcomes (Levine et al. 2020). Offline RL algorithms typically address this with either value conservatism to penalize state-action pairs that don't have data coverage, policy constraints to keep the learned policy within the support of the dataset's 'ss distribution, or ensembles (Kumar et al. 2020b; Fujimoto et al. 2019; Agarwal et al. 2020). There are several standardized benchmark suites for offline RL, including D4RL (Fu et al, 2020 and RI Unplugged (Gülçehre al. 2020), which include primarily deterministic tasks with
8

---

datasets generated by one several suboptimal policies. Implicit Q-Learning (IQL) (Kostrikov ct al. 2021) is similar to SUPERB in that it uses the upper range of a distributional value function and uses it to back out a policy without sampling evaluating out-of- actions. In contrast to IQI. SUPERB is motivated as a data augmentation scheme for RvS agents; further, it does not seek to minimize 1-step Bellman error, approximates the distribution of potential returns rather than an upper expectile of such a distribution, and uses a different approach to policy learning.
RL via Supervised Learning Our work builds off of the paradigm for RL first introduced in Schmidhuber (2019) and Kumar et al. (2019) where a predictive model is trained to predict actions conditioned on a desired outcome such as return. We refer to this class of algorithms as RvS, or RL via supervised learning. following Emmons et al. (2021). Decision Transformer uses large transformer model that conditions on trajectory history and desired returns and is evaluated on offline RL tasks, demonstrating that RvS methods can be competitive in this domain (Chen et al. 2021a)- Others have used RvS to learn goal-conditioned policies, such as in [Ghosh et al. (2021) and Paster al. 2 While RvS methods fall behind other model-frre algorithms in online RL, they are competitive in most popular offline RL benchmarks. In recent work on tuning such methods (Emmons et al. 2021), return- conduiinedd RvS with two-layer MLP was shown to be competitive with Decision Transformer and value-basse methods on many tasks.
TD-Learning Our proposed method takes advantage of the n-step TD relation (Sutton 1988]
Chapter 7). Rather than use dynamic programming to learn a specific policy, however, we capture the data efficiencies of TD learning by explicitly augmenting dataset of Monte Carlo returns (Sutton 1988). Because the distributional value functions learned by SUPERB are initially grounded by supervised targets, this gives it similar flavor to fixed finite horizon TD methods (De Asis al.) 2020), which are resistant to the deadly triad, a common source of training instability for deep RL methods (Van Hasselt al., 2018).
Data augmentation has proven to be a powerful technique for im- proving RL sample efficiency in a variety of common scenarios, domain adaptation and transfer AAndrychowicz et al. 2020), RL from pixels (Yarats et al. RL (Kaelbling) 1993). and object-oriented RL (Pitis et al. 2020). In contrast with the aforemen- tioned methods, which augment agent's observations, our method does data augmentation on returns. This is similar to methods such as ESPER (Paster et al.) 2022). which exchanges return- to-go with learned return label to improve performance in stochastic environments, and RUDDER Aarjona-Medina et al., 2019), which learns a redistributed reward function to improve TD learning in environments with delayed rewards. When simulator is available, human demonstration data may also be augmented "online" via imitation, as done by Mannllearr et al. (2020), who similar to SUPERB, use their imitation policy to stitch together segments from different demonstrations.
Data Augmentation in RL
2020 Laskin et al., (020), goal-driven
5 CONCLUSION
RyS approaches such as Decision Transformer have gained traction recently. but performance on tasks that require the agent to stitch together behaviors from different trajectories has lagged be- hind that of value-based methods (Chen al., 2021al Emmons et al., 2021). We quantify this phenomenon, giving examples of MDPs where the ability to compose behaviors from different tra- jectories gives an exponential improvement trajectory coverage, or outright required to solve the task optimally. To improve RvS, we proposed SUPERB, which augments the returns on which RyS agents are trained with returns computed using an n-step TD relation. On the D4RL benchmark suite, SUPERB dramatically improves the performance of RvS, achieving state-of-the-ee art perfor- mance on half of the AntMaze tasks and improving performance on almost all D4RL-gym tasks.
SUPERB provides a simple and compelling offline data augmentation method for an increasingly important class of RL algorithms. Conveniently, while return augmentation can be applied iter- atively, even one or two iterations of data augmentation with SUPERB can dramatically improve performance, For future work, we plan to investigate further into how the number of data augmen- tation iterations can be tuned in offline manner, well as ways to apply these techniques to non-RL domains that use sequence modeling, such as language modeling. Future work might also consider scaling SUPERB to larger offline datasets in complex environments where coverage gaps are inevitable and data must collected from a wide range of different policies. Finally. it may be interesting to extend SUPERB to other empirical statisties that Rvš methods might condition on.
9

---

REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An Optimistic Perspective Of- fline Reinforcement Learning. In Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 104-114. PMLR, 2020. URUhttp:://pxoeeddie v1 v119/agarwal22c.nnet
mlr. press/
OpenAl: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowiez, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-20) 2020.
Jose A Arjona-Medine Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32, 2019.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint wrXvv:1606.01500. 2016.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Rein- forcement Learning via Sequence Modeling. In Mare' Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Informa- tion Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 15084 15097, 2021a. URL htttps://prooeeddnggg neures wwtpeurisss.cpppppppeeeeo pctppppapee//22//aass... 7F489f (TZ2FFF220000107775555111766666 200777/55311170G333>>ttttttt. html
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 10044-00550597, 2021b.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 32, 2018
Kristopher De Asis, Alan Chan, Silviu Pitis, Richard Sutton, and Daniel Graves. Fixed-horizon temporal difference methods for stable reinforcement learning. In Proceedings of the AAAI Con- ference Artificial Intelligence, volume 34, pp. 3741-3748. 2020.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline via supervised learning? International Conference on Learning Representations, 2021.
Benjamin Eysenbach, Soumith Udatha, Sergey Levine, and Ruslan Salakhutdinov. Imitating past successes can be very suboptimal. CoRR, 007727206.0378, 2022. doi: 10.48550/arXiv.2206. 03378. URL http:httt////l// /ttp::/:///.i..o//........... 485.0/455////0 ar.55///rxiivv.
2206.03378
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference machine learning, pp. 2052-2062. PMLR, 2019.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Ey- senbach, and Sergey Levine. Learning to Reach Goals via Iterated Supervised Learning. In 9th International Conference Learning Representations. ICLR 2021, Virtual Event. Austria, May 3-7, 2021. OpenReview.net, 2021. URL FALAOXO6yNJ
net fformmirdd
10

---

Çaglar Gülgehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel Mankowitz, Cosmin Paduraru, Gabriel Dulac-Aroldd, Jerry Li, Mohammad Norouzi. Matthew Hoffman, Nicolas Heess, and Nando de Freitas. RL Unplugged: A Collection of Benchmarks for Offine Reinforcement Learn- ing. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Ticn Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Con- ference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL httpttsss//prrcoedagg...eeeel
CC wttpp/ppppppp//////ppppppppp
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 1 2021. / http::::c///ccceeeeeeeee hew.:eeeeddiingg...iur.......... new.ips.ccr... pa paper/2021 41t1:///2.1111111044444444444333
Abst .html
Leslie Pack Kaelbling. Learning to achieve goals. In IJCAI, volume 2, pp. 1094-8. Citeseer, 1993.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. International Conference on Learning Representations, 2021.
Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-Conditioned Policies. CoRR, abs/1912. 13465, 2019, URL
13465 arXiv: 1912. 13465.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179 179-1191, 2020a.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-Learning for Of- fline Reinforcement Learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, De- cember 62, 2020, virtual, 2020b. URL httL:ttttpp:s/pcceeeiinnn neurips......................... cttpspccppppppp 2020
32w.22222222122222 3a5.553327111110000002222aaaaaaa Abst htm
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforce- ment learning, pp. 45-73. Springer, 2012.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. Advances in neural information processing systems, 33: 19884-19895, 2020.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers. CoRR, 25552205.15241 2022. doi: 10.48550/arxiw.2225..r5 https::/doi. 48550/04555/200 www.:/arxii.... 2205 15241
URL
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline Reinforcement Learning: Tutorial, Review, and Perspectives Open Problems. arXiv:2005.0164 [cx, staf], November 2020. URL
org/oogg///..o.o wwwgg//bb//....../ 2000.00000000 C1643 arXiv: 2005.01643.
llya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint 1608 (00.::/6..00083,,, 2016.
Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. In 7th International Conference on Learning Reprexentations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
net/forum B-1iidkkgKkGGCC
Ajay Mandlekar, Danfei Xu, Roberto Martm-Marrtn, Silvio Savarese, and Li Fei-Fei. Gti: Leam- ing to generalize across long-horizon tasks from human demonstrations. Robotics: Science and Systems, 2020.
11

---

Keiran Paster, Sheila A. Mcllraith, and Jimmy Ba. Planning from Pixels using Inverse Dynamics Models. 9th International Conference on Learning Representations, ICLR 2021, Virtull Event, Austria, May 3-7, 2021. OpenReview.net, 2021. /
net/ forum?
jd=k 7Ro
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count fail in stochastic enviromments. Advances in Neural Information Processing Systems, 2022.
count on luck: Why decision transformers
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33:3976-3990, 2020.
Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards-juss map them to actions. arXiv preprint arXiv:1912.02875: 2019.
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, (1):9-44, 1988.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Gerald Tesauro. Temporal difference learning and id-gammon, Commun, ACM, 38(3):88668 1995. doi: 10.1 10.1145520333..203344
1ttp:///....11/22033.3 2015///233330.233333333 203343
Hado Van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Mo- dayil. Deep reinforcement learning and the deadly triad. arXiv preprint 1rXiv:1812.0264b 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention All you Need. In Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https::/prooeeedngg.s neurips cc/paperr2017/aaat Abstract
(EEe2221377/611111b3ccccccaaaaa aat html
Adam R Villaflor, Zhe Huang, Swapnil Pande, John M Dolan, and Jeff Schneider. Addressing optimism bias in sequence modeling for reinforcement learning. International Conference on Machine Learning, pp.
PMLR, 2022.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3):279-292, 1992.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representa- fions, 2020.
12

---

A PROOF OF PROPOSITIONS
We prove Proposition 2 before Proposition We use the following Lemma, which is stated two parts for clarity (part (B) generalizes part (A)). The Lemma considers the "N-Step Chain MDP", exemplified in Figure 2 with N - states, and b actions in each state, which each may have different
reward distributions.
Lemma MugmeettaioonBoottttaappa
Equivalence).
(A) For any particular 2-step trajectory, (so, do, To, 81; (1,T1,,,, 82). in a 2-Step Chain MDP, where the subtrajectory тO (80, do, ro, 81) appears n times in a given empirical dataset, and the subtrajectory Ti ($1, 01 ri, 82) appears times, the following return estimators are equivalent:
1. Augmentation:
First form an augmented dataset of nm trajectories, by concatenating each pair of trajectory segments To and т1. Then estimate the return of as the empirical average of its return in the augmented dataset:
G(T)
nm
2. Bootstrapping:
First estimate the return of Ti as G(Ti) Σx ri/m. Then estimate the return of T using the bootstrapped estimator:
G(r) 2i=1 + G(T))
(B) Given an Ststep Chain MDP and an N-step trajectory T, whose 1-step subtrajectories, {Tj = (8j, (.j4)))))))) each appear m2 times in the empirical dataset, the following return estimators are equivalent:
Augmentation:
First form an augmented datasel of Iymi trajectories, by concatenating each combination of trajectory segments. Then estimate the return of T as the empirical average of its return in the augmented dataset:
I.ym Σk r
G(r)
I1my
N steps of Bootstrapping:
Denoting the trajectory slice from 8k to 8N as estimate the return of using the bootstrapped estimator:
G(TN.) E rk + G(Te+1)
mk
where G(TN.) O. Sum of local reward estimators:
Estimate the return of T as a sum of the returns of its 1-step components:
N
G(T)
1=1
where G(TN:) == O.
Proof. Part (A) is special case of Part (B)(1)-(2). Both follow by induction from the trivial base case
Σ m TN-1
MN-1
13

---

30
81
82
N
Figure 4: N-Step Chain MDP: The MDP is initialized in so and terminates in SN+1, with actions always leading the next state in the chain. this MDP there are iN possible trajectories.
where the inductive step is:
my
(r) r 2 1
j-l N-1
G(Tr.)
1=1
k=
1
IT: -1 j=b
my
my j=(+1 my
N- my 5 N-1
j=z m;
Liml
2==lt1 rk
11
m mIIl
my m m-
j=
Di ri m,G(Te+1:)
m
my
Σ me
ii r + G(Te+1))
m
Part (B)() follows by repeatedly unrolling the last term in equation above.
As straightforward consequence of above Lemma, we have:
Proposition 2 (Coverage of Trajectory Space). deterministic N-step Chain MDP with branch- ing factor b, it is possible for N -seep bootstrapping (implicitly) capture full' "coverage" of the size & trajectory space with only empirical trajectories-an exponential increase in coverage relative bootstrapping. Coverage here refers the percentage of unique trajectories present in the dataset, where we consider two trajectories with the same actions taken in every state as equivalent.
Proof. This follows directly from Lemma B))))-(2 if every possible 1-step subtrajectory ( (s,, aj, $f+1)) appears exactly once in the b empirical trajectories.
We can additionally make the following related statement about the sample complexity of valuing each trajectory in the Chain MDP when rewards are stochastic (note that in the Chain MDP, trajectory is equivalent bound on "every policy"" policy valuation in the Chain MDP).
policy, so that the following proposition
finite sample complexity
Proposition 2a (Sample Complexity of Policy Valuation in Stochastic Case). Consider N N-Step Chain MDP with branching factor b, whose rewards at each action are bounded random variables with R(s, a) [0, 1].
(A) Suppose we have n samples )= of each length N trajectory. There are b such trajectories, providing total of N samples of length subtrajectories. Without bootstrapping, if
N2
N
N
2
log
8 or equivalently, m
log
22
2.2
8
then, with probability at least - 8, we have:
Σ. G(T(0) E [G((T)]
max.
Sc.
14

---

(B) Suppose have least & samples of each length 1 subtrajectory T; (sj, aj, 8(+1).
There are bN such subtrajectories, providing total of m = b samples of length subtrajectories. Using N N-Step bootstrapping described in Lemma 4 if
N log
2bN
b log
2bN
6 or equivalently, m
2ea
then, with probability at least
ô, we have:
maz+ E, G(rl0) - BB(G(T) cc
Therefore, in an N-Step Chain MDP with > 1, plain RyS requires NN1 times as many (i.e.,
exponentially more) samples obtain the same precision N-step bootstrapping.
Proof.
(A) This follows from Hoeffding's Inequality by noting that G(T) E [0, _V] and taking union bound over the b- trajectories.
(B) This follows from Hoeffding's Inequality using e' e/N, so that the trajectory error (a sum of N 1-step subtrajectories) is bounded by c, and taking union bound over the b.N 1-step subtrajectories.
The Chain MDP used in Proposition revealss a general construction to prove necessity, as follows.
Proposition 1 (Necessity). For any positive integer n, there exists an MDP M and data generating policy Te for which n-step bootstrapping is strictly necessary generate dataset D containing (s, (s.0,9) 9"), where g" is the maximum reward achievable for a trajectory starting in (s,a).
Proof. For the general case, consider an state Chain MDP with single optimal trajectory, where 'D is generated by a combination of n non-Markovian policies that each contain a single, unique length 1 segment of the optimal trajectory. single step of SUPERB bootstrapping can add at most I new length I segment to any particular (augmented) return label, so that n steps of SUPERB bootstrapping are necessary to compose all length 1 segments from the optimal policy.
B IMPLEMENTATION DETAILS
Code for our implementation and experiments is available at [GITHUB LINK TO BE ADDED). We implemented our experiments top of the code for RvS (Emmons al. 2021) found
https://qithub. com/ www.tteettooo/rrrroo
B. RETURN MODEL
To model the distribution of returns we use an ensembled Quantile Regression Network (QRN), as proposed by Dabnnyy al.) (2018), which maps states value distribution represented by 20 quantiles. The quantile regression is trained using the Huber quantile loss proposed by Dabney et al. 2018) with & = 1, and is optimized for epochs (Antmaze) or epochs (Gym) using Adam Loshchilov & Hutter 2019) using batch size of 1024, and constant learning rate of le-3 for AntMaze and 3e-4 for the Gym tasks.
Our QRN is ensemble of 5 feedforward neural networks, each with 3 hidden layers of 512 neurons and ReLU activations. Both the inputs to the networks, and the output targets are normalized. For AntMaze experiments, apply value clipping clip target return values to their known feasible range (after reward transformation) of [-1/(1 y), Oj. where & 0.997. In each bootstrapping step, new randomly initialized QRN is trained the current augmented dataset.
15

---

To form augmented labels, first use the QRN to propose return labels for all trajectory suffixes in the dataset. This is done by calling each member the QRN ensemble on the first observation in the trajectory suffix, taking the mean (Antmaze) minimum (Gym) across the ensemble, and average the top quantiles of the result.
Proposed return labels in hand, then apply the following backward induction procedure relabel all return labels D:
def e__augmented. ceremenned_return_ceo lawerrnnaaalllleeeeeeeeeelelllle (traj, proposed labels, rewards traj.rewards returns ret traj_len - len (rewards)
discount_ factor):
for reversed (range (traj_len)): factor
ret fastor ret += (float (rewards(il))
if (i + +1 1 traj_len): ret max (ret, float (rewards +\ discount_factor proposed 11) returns. append (ret)
12
13
14
returns list (reversed (returns))
return returns
On each iteration of SUPERB, we keep only the most recently generated augmented labels, and discard both the QRN for that step and the data it was trained on. After the final step, we train one more QRN use as the value function for the purpose of forming reward targets for the Rv$ policy, described next.
B.2 RvS POLICY
For simplicity, we base our policy learning heavily off of the hyperparameters discussed in Emmons et al. (202T) As in Emmons et al. (2021), our policy consists of simple feedforward neural network with two hidden layers with width equal to 1024. We optimized our policy network using Adam W [oshchilov & Hutter 2019). applying weight decay of 1e-2, learning rate of 1e-3. and batch size of 2048, The learning rate is annealed using a cosine learning rate schedule Loshchilov &
Hutter 2016).
In order to support conditioning on higher values of return, we implemented input normalization for the policy network, where inputs are normalized by the mean and standard deviation of the inputs from the dataset. Due to this normalization, policy training could be completed in fewer epochs and we decreased the number of training epochs from 2000 400 for D4RL-gym and trained for only 100 epochs AntMaze.
As discussed section 3.5 we dynamically choose a return target g" by using our learned return model (for conditioning we use only the top quantile, rather than take the mean of the top 5, as done for augmentation). However, we found that we can additionally improve performance by increasing the target by some A to just above the value predicted by (g s). We tune A separately for each environment and number of iterations, as the return distribution varies with these variables. Table shows that still obtain good performance using fixed A across all tested environments.
16

---

Tisk
RvS moo
SB (1)
SB (2)
SB (3)
Rv8 Ioei Fixod à
SB (1) Fixol
SB
2D--l &
SB (3) Fixod A
antmaze umaze-v2
84.2
85.6
$5.4
83.8
79.2
8.2
0.0
0.0
diverse
71.7
76.5
$1.2
76.9
76.5
76.2
75.4
medium play
antmnre mudium--ineeee--
2:7
78.1
66.2
69.6
48.1
intmane- mntmire-arge Sarge-plag--
0.8
16.2
25.8
26.5
D.D
15.8
21.5
24.2
antmare divereee----2222
6.9
24.2
36.9
19.6
4,2
18.1
36.9
19.6
halfcheetah candom-v2
2.2
2.2
2.2
2.3
2.2
2.2
2.2
Hupper-radom--e
5.8
7.9
21.9
2illee22--mnoom-22
6.0
6.3
5.6
5.6
39.1
39.
39.1
39.7
38.3
38.7
38.1
39.5
penper-mddium-rppie repply------
77.9
9016
90.3
91.0
76.3
98.2
88.8
89.7
50.5
59.9
medium-x2
hopper-modiue meduu--ddiun-222
58.9
57.6
56.8
57.1
56.2
56.4
54.9
wa walker2d (arr22--mddiu---
72.4
73.5
74.7
74.2
72.4
72.1
73.4
73.8
hallchoctall- modium-cxppet-122222
92.7
92.5
92.5
92.7
92.0
92.3
92.4
92.3
hoppes-mediime expert--22
108.5
109.5
109.8
Seneeerdd-ddiubseeo expert-v2
106.9
107.9
107.6
108.0
Table 3: Results when using the best fixed A across all D4RL environments, Although there is some advantage to tuning A for each environment, we can obtain good performance using fixed A.
Task
CQL
IQL
BC
% BC
DT
RyS-R
RvS
SB (1)
SB (2)
SB (3)
antmare-umaze-vv
44.8
83.7
54.6
60.0
65.6
64.4
84.2
85.6
85.4
83.8
antmaze-umaz-ddiveneee
23.4
66.7
51.2
70.1
77.7
commare-medium saaiiiuumplly--2
0.0
74.7
0.0
42.1
4.5
2.7
58.1
76.5
66.5
cedmmaze-mediunddivee
0.0
71.0
0.0
37.2
0.6
7.7
2.7
55.8
78.1
66.2
intmaze-large-playy-t
0.0
44.3
0.0
28.0
3,5
16.2
26.5
antmaze-large diverse-v2
0.0
49.0
0.0
34,3
0.2
3.7
6.9
24.2
36.9
19.6
halfchectah raadoom---
18.6
2.3
2.0
2.2
3.9
2.2
2.2
2.2
2.3
hopper-random-22
9.3
4.8
4.1
7.5
7.7
5.8
7.9
14.3
23.4
walker2d-randdm-y random-v2 walker2d
2.5
1.7
1.7
-0.2
5.6
halfcheetah wew.ataaa-mmiiuu-rrerra--.
44.2
36.6
40.6
36.6
39.1
39.1
39.7
97.8
94.7
18.1
75.9
82.7
73.5
77.9
90.6
90.3
91.0
walker2d meadium--replay--
86.1
73.9
26.0
62.5
66.6
60.6
50.5
56.5
66.2
59.9
halfchoctah- haffchettal-ncdiun-22
49.1
47.4
42.6
42.5
42.6
41.6
42,3
42.3
42.1
42.4
64.6
66.3
56.9
67.6
60.2
58.9
57.6
56.8
walker2d-mediune
82.9
78..3
75.3
75.0
74.0
71.7
72.4
73.5
74.7
74.2
halfcheetah medium-expentt experttt22
85.8
86.7
55.2
92.9
86.8
92.2
92.7
92.5
92.5
92.7
hopper-mediuu-exppeee walker2d medium- experttv22
102.0
91.5
52.5
110.9
107.6
101.7
109.8
109.9
110.1
108.5
109.5
109.6
107.5
109.0
108.1
106.0
106.9
108.0
107.9
107.6
Table 4: evaluated our method with bootstrap iterations, denoted SB(1), SB(2), and SB(3) on the AntMaze and Gym tasks from D4RL (Fu et al., 2020). Across the AntMaze tasks, which require stitching together behaviors from several suboptimal trajectories in order to perform well, SUPERB provides strong performance improvement bootstrapping, outperforming the prior state-of-the-art in 3 of the mazes. The benefits of SUPERB on the Gym tasks are more modest, though we observe improvements over RyS on a few of the suboptimal datasets, particularly those where the dataset was generated by a non-Markovian com- bination of policies. Generally, most of the performance benefit seems to be captured by a single iteration of S.PERB.
the baseline RvS agents that don't 't use
Results averaged over
17

---

70
Rvs
60
Pafemaace Improvennnt
20
30
20
10-
antuaae nimaae divense 1/
maller26rroxdoms rod--oodoom- n2
halldhetth medinnm expert -22
laall beettt www.innnnm--- vi
hallehettil ane moddim dinn cop.lppppllee e2
holdhontall- coadom-ve 12
wai Mortal wow-ooo--taa commonee 22
copppeee weseooooooo 22
antiasss horge Daop 22
22
mossee mnsti lime play -2
hosper- medd-n-dnnnn
hoggurt -t dianneee wn
ant intionaee
Figure 5: SUPERB gives a substantial boost in performance compared to RvS (ours) across most tasks. This includes both D4RL AntMaze and Gym tasks.
18


Reviews:
[{"review_id": "_-irXKgevC", "reviewer": "Reviewer_JPnx", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a data augmentation method for return-conditioned supervised learning that takes into account temporal compositionality. The proposed method shows higher performance compared to vanilla baselines.", "summary_of_the_review": "The paper provides a data augmentation method for return-based supervised learning in offline RL. While finding an appropriate data augmentation technique for return-based supervised learning (and in general, offline RL) is important, the proposed method seems to have some drawbacks.", "strength_and_weaknesses": "Strengths\n- The paper deals with an important problem of return-conditioned supervised learning methods, which is that they cannot stitch trajectories to find optimal behavior.\n- The proposed return augmentation method does seems to improve performance on various benchmarks.\n\nWeaknesses\n- The source code is not provided, which makes the credibility of the work questionable.\n- As the authors note, a major advantage of the return-conditioned SL methods is that they do not require value bootstrapping and thus does not fall into the infamous 'deadly triad'. However, adopting the return augmentation with iterative value updates can reintroduce this bootstrapping problem.\n- The return conditioning procedure is not explained in detail even though the algorithm's performance depends heavily on it (Table 2). How is the value model V_\\phi used to sample \"high, but plausible returns for the current episode\"? Also, Appendix B.2 notes that the targeted value generated from the value model is increased by some \\Delta. How is this \\Delta tuned? Overall, how much does each design choice affect the performance?\n- The performance of the proposed method is highly sensitive to N, the number of bootstrapping iterations, especially on Antmaze.\n- The D4RL benchmark results are averaged over only 3 seeds. Adding at least 1~2 more seeds would be more credible.\n\nQuestions\n- What does the numbers on Figure 3 (b) mean? Does it refer to the expected value (from V_\\phi) ?\n\n############## Post-rebuttal comment ##############\n\nI appreciate the authors for updating the manuscript, but there still are some remaining concerns.\n\nFirst, the hyperparameter tuning procedure is still somewhat ambiguous. The tuning process (search range, selection method, ...) for \\Delta is not included in the paper. Also, the authors mention that \"each method has one set of hyperparameters shared across all tasks\" on the updated Table 1, but this is wrong since the proposed method tunes \\Delta per task.\n\nSecond, I agree that the work can be interpreted as unifying value-based methods and supervised-RL, but given the empirical results, I think the proposed method fails to achieve the best of both worlds (the proposed method outperforms value-based baselines on only Antmaze). I encourage the authors to improve the method so that the cons from both approaches are minimized while the pros are maximized.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Witr_HYU29", "reviewer": "Reviewer_FZYT", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors observed that RvS, their main baseline method, does not do well on Antmaze when not given privileged information such as the goal location. They hypothesized that by modifying the reward to do bootstrapping, similar to other methods based on temporal difference learning(e.g. CQL, BCQ), that will allow the RvS methods to do better. \nThey observed no significant benefit on halfcheetah, hopper, and walker2d tasks (apart from the random data split which benefits more from trajectory stitching), but they outperformed the previous SOTA in 3 of the 6 mazes in AntMaze.\n", "summary_of_the_review": "This paper proposes SuperB, an incremental improvement over RvS where they introduced bootstrapping to the target reward. The method is shown to work well where it is known that the offline data contains multiple imperfect trajectories that can be combined together to form a better one. I recommended a rating of 6 due to their sound experimental results and due to the limited novelty, but also due to the paper’s thorough discussion of the method’s limitations and future works. ", "strength_and_weaknesses": "\nThe paper has a clear motivation and diagnosis for the problem they are trying to tackle, a sound logic for their proposed solution (SuperB) accompanied by a proof. They showed that their proposed method improves over the baseline RvS method (the one without privileged goal info) by a significant margin. \n\nThe main weakness is where else this approach can be used. As shown in Table 1, it does not result in SOTA performance unless it is clear from the dataset and problem setup that an agent will benefit from stitching together multiple trajectories.  It is also assumed that different trajectories will visit similar states, allowing the bootstrapping to be effective. SuperB is designed specifically for those assumptions and works well on AntMaze medium and umaze. It is less clear why the method works less well on antmaze large compared to IQL.\nThe authors are aware of the issues of their proposed method and as mentioned in 3.4, “there is no universally optimal number of iterations”, meaning that one needs to either use domain knowledge (like in RvS-G) or do large scale hyperparameter sweeping to find the optimal setup. \nNevertheless, the method looks like a strict improvement over the baseline RvS on the D4RL dataset, and assuming reasonably that the state representation is reasonably good, there does not seem to be a strong downside to swapping out RvS rewards for the SuperB rewards. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ttc69Oe_j1U", "reviewer": "Reviewer_sHVj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposed using value bootstrapping to augment the return data for learning decision transformers or RvS (RL via supervised learning) in general. This approach, which performs explicit bootstrapping / stitching, allows mapping optimal actions to optimal return-to-gos that do not appear in the actual experience and thus achieve better performance. ", "summary_of_the_review": "I lean to acceptance. However, I would like to see if the authors can improve on clarity of implementation details, provide more insight from sensitivity / ablation studies, and increase reproducibility.", "strength_and_weaknesses": "Strength:\n* Very interesting and novel proposal that addresses the critical bottleneck of decision transformers (DT) or RvS of not being able to infer the optimal return by bootstrapping. I am not aware of previous works that successfully combine DT/RvS and TD-bootstrapping.\n* The empirical results on D4RL are very impressive, especially on AntMaze, where an optimal policy must stitch together behaviors from several suboptimal demonstrations, from almost zero to reasonable performance.\n\nWeakness\n* How quantile regression was perform and how bias was added into sampling (specifically the implementation related to Dabney et al. (quantile regression) and Lee et al. (multi-game decision transformers)) was not clearly described. I suppose these could be very important to the stability of offline TD-learning. I would like to see more details and potentially more sensitivity analysis on these. \n* Many details are missing: Hyper-parameters related to the above point are missing. Network architecture is also missing. Although the reader can refer to code from previous work, it's better to make the paper self-contained and improve reproducibility.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-GOkSgiVmZd", "reviewer": "Reviewer_RPpv", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The submission introduces SuperB, a new approach for offline RL following the RL via Supervised Learning (RvS) framework. SuperB addresses a shortcoming of RvS methods, which fail to stitch together trajectories that were executed as parts of non-optimal trajectories in the data, but which would result in an optimal (or high-performing) trajectory if combined into a single trajectory. The approach is simple conceptually: train a return model to estimate the distribution of returns from any given state, and use that to compute potentially higher returns for states in the data set that would be attained if a different trajectory had been followed thereafter, thereby augmenting the data with better trajectories. The authors show theoretically a necessity result (that n-steps returns are necessary for the data set to contain the optimal return if the data-generating policy is Markovian) and a coverage result (that n-step returns enable generating up to exponentially many unseen trajectories). Empirically, SuperB performs competitively to value-based RL and outperforms RvS methods in settings where the data-generating policy is non-Markovian. \n", "summary_of_the_review": "Overall, I think this submission is a good contribution and I'm leaning towards recommending its acceptance. My main concerns relate to the motivation for the proposed approach, the fairness of the experimental evaluation, and the somewhat overstated theoretical results. I would like to see the authors' response to these and would be glad to revise my recommendation accordingly.\n\n############# Update after rebuttal #############\n\nI am updating my score from 6 (marginally above threshold) to 5 (marginally below threshold) per the discussion with the authors below.", "strength_and_weaknesses": "########### Strengths ###########\n1. The paper is clearly written and easy to follow. The motivation is specified clearly from the beginning and repeated/clarified throughout the paper\n2. The proposed SuperB approach is simple conceptually, and is a reasonable solution to handle the problem of trajectory stitching from offline data, as demonstrated theoretically and empirically\n3. The idea of using a return model to stitch trajectories, instead of attempting to explicitly find trajectories in the data that could be stitched, is clever\n\n########### Weaknesses ###########\n1. The theoretical results might be somewhat overstated, and they appear to fall out by construction. They do provide some grounding to validate the proposed solution, so I do think they should be kept in the paper, if perhaps rephrased and downplayed\n2. The resulting approach is not particularly strong compared to baselines, especially value-based approaches. While state-of-the-art performance shouldn't be the focus of publications, it is important to discuss and analyze where these ideas might take us in the future if pursued further\n3. Some details about the experimental setting are missing and would be necessary for ensuring a fair comparison\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject