SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anirudhan Badrinath
Yannis Flet-Berliac
Allen Nie
{abadrina, yfletberliac, anie, ebrun}@cs.stanford.edu
Despite the recent advancements in offline reinforcement learning via supervised learning (RvS) and the success of the decision transformer (DT) architecture in various domains, DTs have fallen short in several challenging benchmarks. The root cause of this underperformance lies in their inability to seamlessly connect segments of suboptimal trajectories. To overcome this limitation, we present a novel approach to enhance RvS methods by integrating intermediate targets. We introduce the Waypoint Transformer (WT), using an architecture that builds upon the DT framework and conditioned on automatically-generated waypoints. The results show a significant increase in the final return compared to existing RvS methods, with performance on par or greater than existing popular temporal difference learning-based methods. Additionally, the performance and stability improvements are largest in the most challenging environments and data configurations, including AntMaze Large Play/Diverse and Kitchen Mixed/Partial.
Traditionally, offline reinforcement learning (RL) methods that compete with state-of-the-art (SOTA) algorithms have relied on objectives encouraging pessimism in combination with value-based methods. Notable examples of this approach include Batch Conservative Q-Learning (BCQ), Conservative Q-Learning (CQL), and Pessimistic Q-Learning (PQL) [Fujimoto et al., 2019, Kumar et al., 2020, Liu et al., 2020]. However, these methods can be challenging to train and often require intricate hyperparameter tuning and various tricks to ensure stability and optimal performance across tasks.
Reinforcement learning via supervised learning (RvS) has emerged as a simpler alternative to traditional offline RL methods [Emmons et al., 2021]. RvS approaches are based on behavioral cloning (BC), either conditional or non-conditional, to train a policy. Importantly, these methods eliminate the need for any temporal-difference (TD) learning, such as fitted value or action-value functions. This results in a simpler algorithmic framework based on supervised learning, allowing for progress in offline RL to build upon work in supervised learning. There are several successful applications of RvS methods, including methods conditioned on goals and returns [Kumar et al., 2019, Janner et al., 2021, Ding et al., 2019, Chen et al., 2021, Emmons et al., 2021].
However, RvS methods have typically struggled in tasks where seamlessly connecting (or "stitching") appropriate segments of suboptimal training trajectories is critical for success [Kumar et al., 2022]. For example, when tasked with reaching specific locations in the AntMaze maze navigation environment or completing a series of tasks in the FrankaKitchen environment, RvS methods typically perform significantly worse than TD learning methods such as Implicit Q-Learning [Fu et al., 2020, Kostrikov et al., 2021].
In this study, we leverage the transformer architecture [Vaswani et al., 2017] to construct an RvS method. As introduced by Chen et al. [2021], the decision transformer (DT) can perform conditional behavioral cloning in the context of offline RL. However, similar to other RvS methods, DT proves inferior in performance across popular Gym-MuJoCo benchmarks compared to other value-based offline RL methods, with a  $15\%$  relative reduction in average return and lowered stability (Table 1).
To tackle these limitations of existing RvS methods, we introduce a waypoint generation technique that produces intermediate goals and more stable, proxy rewards, which serve as guidance to steer a policy to desirable outcomes. By conditioning a transformer-based RvS method on these generated targets, we obtain a trained policy that learns to follow them, leading to improved performance and stability compared to prior offline RL methods. The highlights of our proposed approach are as follows:
- We propose a novel RvS method, Waypoint Transformer, using waypoint generation networks and establish new state-of-the-art performance, in challenging tasks such as AntMaze Large and Kitchen Partial/Mixed [Fu et al., 2020] (Table 1). On tasks from Gym-MuJoCo, our method rivals the performance of TD learning-based methods such as Implicit Q-Learning and Conservative Q-Learning [Kostrikov et al., 2021, Kumar et al., 2020], and improve over existing RvS methods.
- We motivate the benefit of conditioning RvS on intermediate targets using a chain-MDP example and an empirical analysis of maze navigation tasks. By providing such additional guidance on suboptimal datasets, we show that a policy optimized with a behavioral cloning objective chooses more optimal actions compared to conditioning on fixed targets (as in Chen et al. [2021], Emmons et al. [2021]), facilitating improved stitching capability.
- Our work also provides practical insights for improving RvS, such as significantly reducing training time, solving the hyperparameter tuning challenge in RvS posed by Emmons et al. [2021], and notably improved stability in performance across runs.
Many recent offline RL methods have used fitted value or action-value functions [Liu et al., 2020, Fujimoto et al., 2019, Kostrikov et al., 2021, Kumar et al., 2020, Kidambi et al., 2020, Lyu et al., 2022] or model-based approaches leveraging estimation of dynamics [Kidambi et al., 2020, Yu et al., 2020, Argenson and Dulac-Arnold, 2020, Shen et al., 2021, Rigter et al., 2022, Zhan et al., 2021].
RvS, as introduced in Emmons et al. [2021], avoids fitting value functions and instead leverages behavioral cloning. In many RvS-style methods, the conditioning variable for the policy is based on the return [Kumar et al., 2019, Srivastava et al., 2019, Schmidhuber, 2019, Chen et al., 2021], but other methods use goal-conditioning [Nair et al., 2018, Emmons et al., 2021, Ding et al., 2019, Ghosh et al., 2019] or leverage inverse RL [Eysenbach et al., 2020]. Recent work by Brandonbrener et al. [2022] has explored the limitations of reward conditioning in RvS. In this study, we consider both reward and goal-conditioning.
Transformers have demonstrated the ability to generalize to a vast array of tasks, such as language modeling, image generation, and representation learning [Vaswani et al., 2017, Devlin et al., 2018, He et al., 2022, Parmar et al., 2018]. In the context of offline RL, decision transformers (DT) leverage a causal transformer architecture to fit a reward-conditioned policy [Chen et al., 2021]. Similarly, [Janner et al., 2021] frame offline RL as a sequence modeling problem and introduce the Trajectory Transformer, a model-based offline RL approach that uses the transformer architecture.
Algorithms building upon the DT, such as online DT [Zheng et al., 2022], prompt DT [Xu et al., 2022] and Q-Learning DT [Yamagata et al., 2022], have extended the scope of DT's usage. Furuta et al. [2021] introduce a framework for hindsight information matching algorithms to unify several hindsight-based algorithms, such as Hindsight Experience Replay [Andrychowicz et al., 2017], DT, TT, and our proposed method.
Some critical issues with DT unresolved by existing work are (a) its instability (i.e., large variability across initialization seeds) for some tasks in the offline setting (Table 1) and (b) its relatively poor performance on some tasks due to an inability to stitch segments of suboptimal trajectories [Kumar et al., 2022] – in such settings, RvS methods are outperformed by value-based methods such as
Implicit Q-Learning (IQL) [Kostrikov et al., 2021] and Conservative Q-Learning (CQL) [Kumar et al., 2020]. We address both these concerns with our proposed approach, demonstrating notably improved performance and reduced variance across seeds for tasks compared to DT and prior RvS methods (Table 1).
One of the areas of further research in RvS, per Emmons et al. [2021], is to address the complex and unreliable process of tuning hyperparameters, as studied in Zhang and Jiang [2021] and Nie et al. [2022]. We demonstrate that our method displays low sensitivity to changes in hyperparameters compared to Emmons et al. [2021] (Table 2). Further, all experiments involving our proposed method use the same set of hyperparameters in achieving SOTA performance across many tasks (Table 1).
We assume that there exists an agent interacting with a Markov decision process (MDP) with states  $s_t \in S$  and actions  $a_t \in \mathcal{A}$  with unknown transition dynamics  $p(s_{t+1} \mid s_t, a_t)$  and initial state distribution  $p(s_0)$ . The agent chooses an action sampled from a transformer policy  $a_t \sim \pi_\theta(a_t \mid s_{t-k\dots t}, \Phi_{t-k\dots t})$ , parameterized by  $\theta$  and conditioned on the known history of states  $s_{t-k\dots t}$  and a conditioning variable  $\Phi_{t-k\dots t}$ . Compared to the standard RL framework, where the policy is modeled by  $\pi(a_t \mid s_t)$ , we leverage a policy that considers past states within a fixed context window  $k$ .
The conditioning variable  $\Phi_t$  is a specification of a goal or reward based on a target outcome  $\omega$ . At training time,  $\omega$  is sampled from the data, as presented in Emmons et al. [2021]. At test time, we assume that we can either generate or are provided global goal information  $\omega$  for goal-conditioned tasks. For reward-conditioned tasks, we specify a target return  $\omega$ , following Chen et al. [2021].
To train our RvS-based algorithm on an offline dataset  $\mathcal{D}$  consisting of trajectories  $\tau$  with conditioning variable  $\Phi_t$ , we compute the output of an autoregressive transformer model based on past and current states and conditioning variable provided in each trajectory. Using a negative log-likelihood loss, we use gradient descent to update the policy  $\pi_{\theta}$ . This procedure is summarized in Algorithm 1.
Algorithm 1 Training algorithm for transformer-based policy trained on offline dataset  $\mathcal{D}$
Input: Training dataset  $\mathcal{D} = \{\tau_1,\tau_2,\tau_3,\dots,\tau_n\}$  of training trajectories. for each  $\tau = (s_0,a_0,\Phi_0,s_1,a_1,\Phi_1,\ldots)$  in  $\mathcal{D}$  do Compute  $\pi_{\theta}(a_t\mid s_{t - k\cdot t},\Phi_{t - k\cdot t})$  for all  $t$  Calculate  $L_{\theta}(\tau) = -\sum_{t}\log \pi_{\theta}(a_t\mid s_{t - k\cdot t},\Phi_{t - k\cdot t})$  Backpropagate gradients w.r.t.  $L_{\theta}(\tau)$  to update model parameters end for
4 Waypoint Generation
In this section, we propose using intermediate targets (or waypoints) as conditioning variables as an alternative to fixed targets, proposed in Emmons et al. [2021]. Below, we motivate the necessity for waypoints in RvS and present a practical technique to generate waypoints that can be used for goal-conditioned (Section 4.2) and reward-conditioned tasks (Section 4.3) respectively.
4.1 Illustrative Example
To motivate the benefits of using waypoints, we consider an infinite-horizon, deterministic MDP with  $H + 1$  states and two possible actions at non-terminal states. A graphical representation of the MDP is shown in Figure 1. For this scenario, we consider the goal-conditioned setting where the target goal state during train and test time is  $\omega = s^{(H)}$ , and the episode terminates once
Figure 1: Chain MDP to motivate the benefit of intermediate goals for conditional BC-based policy training.
In offline RL, the data is often suboptimal for achieving the desired goal during testing. In this example, suppose we have access to a dataset  $\mathcal{D}$  that contains an infinite number of trajectories
collected by a random behavioral policy  $\pi_{b}$  where  $\pi_{b}(a_{t} = a^{(1)}\mid s_{t}) = \lambda >0$  for all  $s_t$ . Clearly,  $\pi_{b}$  is suboptimal with respect to reaching  $\omega$  in the least number of timesteps; in expectation, it takes  $\frac{H}{1 - \lambda}$  timesteps to reach  $s^{(H)}$  instead of  $H$  (optimal) since the agent "stalls" at the current state with probability  $\lambda$  and moves to the next state with probability  $1 - \lambda$ .
Consider a global goal-conditioned policy  $\pi_G(a_t\mid s_t,\omega)$  that is optimized using a behavioral cloning objective on  $\mathcal{D}$ . Clearly, the optimal policy  $\pi_G^* (a_t\mid s_t,\omega) = \pi_b(a_t\mid s_t)\forall s_t$  since  $\omega = s^{(H)}$  is a constant. Hence, the global goal-conditioned policy  $\pi_G^*$  is as suboptimal as the behavioral policy  $\pi_{b}$
Instead, suppose that we condition a policy  $\pi_W(a_t\mid s_t,\Phi_t)$  on an intermediate goal state  $\Phi_t = s_{t + K}$  for some chosen  $K < \frac{1}{1 - \lambda}$  (expected timesteps before  $\pi_b$  executes  $a_2$ ), optimized using a behavioral cloning objective on  $\mathcal{D}$ . For simplicity, suppose our target intermediate goal state  $\Phi_t$  for some current state  $s_t = s^{(h)}$  is simply the next state  $\Phi_t = s^{(h + 1)}$ . Based on data  $\mathcal{D}$  from  $\pi_b$ , the probability of taking action  $a^{(2)}$  conditioned on the chosen  $\Phi_t$  and  $s_t$  is estimated as:
$$
\begin{array}{l} \operatorname * {P r} _ {\pi_ {b}} [ a _ {t} = a ^ {(2)} \mid s _ {t} = s ^ {(h)}, s _ {t + K} = s ^ {(h + 1)} ] = \frac {\operatorname * {P r} _ {\pi_ {b}} [ a _ {t} = a ^ {(2)} , s _ {t + K} = s ^ {(h + 1)} \mid s _ {t} = s ^ {(h)} ]}{\operatorname * {P r} _ {\pi_ {b}} [ s _ {t + K} = s ^ {(h + 1)} \mid s _ {t} = s ^ {(h)} ]} \\ = \frac {(1 - \lambda) \lambda^ {K - 1}}{\binom {K} {1} [ (1 - \lambda) \lambda^ {K - 1} ]} = \frac {1}{\binom {K} {1}} = \frac {1}{K}. \\ \end{array}
$$
Hence, for the optimal intermediate goal-conditioned policy  $\pi_W^*$  trained on  $\mathcal{D}$ , the probability of choosing the optimal action  $a^{(2)}$  is:
$$
\pi_ {W} ^ {*} \left(a _ {t} = a ^ {(2)} \mid s _ {t} = s ^ {(h)}, \Phi_ {t} = s ^ {(h + 1)}\right) = \frac {1}{K}.
$$
Since  $\pi_G^* (a_t = a^{(2)}\mid s_t = s^{(h)},\omega) = 1 - \lambda$  and we choose  $K$  such that  $\frac{1}{K} >1 - \lambda$  , we conclude:
$$
\pi_ {W} ^ {*} (a _ {t} = a ^ {(2)} \mid s _ {t}, \Phi_ {t}) > \pi_ {G} ^ {*} (a _ {t} = a ^ {(2)} \mid s _ {t}, \omega).
$$
The complete derivation is presented in Appendix A. Based on this example, conditioning the actions on reaching a desirable intermediate state is more likely to result in taking the optimal action compared to a global goal-conditioned policy. Effectively, the conditioning acts as a "guide" for the policy, directing it toward desirable intermediate targets in order to reach the global goal.
4.2 Intermediate Goal Generation for Spatial Compositionality
In this section, we address RvS's inability to "stitch" subsequences of suboptimal trajectories in order to achieve optimal behaviour, based on analyses in Kumar et al. [2022]. In that pursuit, we introduce a technique to generate effective intermediate targets to better facilitate stitching and to guide the policy towards desirable outcomes, focusing primarily on goal-conditioned tasks.
Critically, the ability to stitch requires considering experiences that are more relevant to achieving appropriate short-term goals before reaching the global goal. To illustrate this, we show a maze navigation task from the AntMaze Large environment in Figure 2, where the evaluation objective is to reach a target location from the start location [Fu et al., 2020].
Analyzing the training trajectories that pass through either the start location (blue) or target location (red), less than  $5\%$  of trajectories extend beyond the stitching region into the other region, i.e., the target or start regions respectively. Since trajectories seldom pass through both the start and target regions, the policy must "stitch" together subsequences from the blue and red trajectories within the stitching region, where the trajectories overlap the most. By providing intermediate targets within this region, rather than conditioning solely on the global goal, we can guide the policy to connect the relevant subsequences needed to reach the target effectively.
Figure 2: antmaze-large-play-v2 task to navigate from the start location (circle) to the target location (star). Blue and red colored lines are training trajectories passing through the start or end locations respectively.
To obtain effective intermediate targets, we propose the goal waypoint network, explicitly designed to generate short-term, intermediate goals. Similar to the illustrative example in Section 4.1, the purpose of these intermediate targets is to guide the policy network  $\pi_{\theta}$  towards states that lead to the desired global goal by facilitating stitching of relevant subsequences.
To that end, we represent the goal waypoint network  $W_{\phi}$ , parameterized by  $\phi$ , as a neural network that makes approximate  $K$ -step predictions of future observations conditioned on the current state,  $s_t$ , and the target goal,  $\omega$ . Formally, we attempt to minimize the objective in Equation 1 across the same offline dataset  $\mathcal{D}$ , where  $L_{\phi}$  is a mean-squared error (MSE) loss for continuous state spaces:
$$
\arg \min  _ {\phi} \sum_ {\tau \in \mathcal {D}} L _ {\phi} \left(W _ {\phi} \left(s _ {t}, \omega\right), s _ {t + K}\right). \tag {1}
$$
While our approach to intermediate target generation seems simple in relation to the complex problem of modeling both the behavioral policy and transition dynamics, our goal is to provide approximate short-term goals to facilitate the downstream task of reaching the global goal  $\omega$ , rather than achieving perfect predictions of future states under the behavioral policy.
4.3 Proxy Reward Generation for Bias-Variance Reduction
In this section, we address the high bias and variance of conditioning variables used by prior RvS methods in reward-conditioned tasks, such as Emmons et al. [2021] and Chen et al. [2021]. Analogously to Section 4.2 (i.e., for goal-conditioned tasks), we propose a technique to generate intermediate reward targets for reward-conditioned tasks to mitigate these issues.
Existing methods rely on either an initial cumulative reward-to-go (desired return) or an average reward-to-go target, denoted as  $\omega$ . Importantly, the former is updated using rewards obtained during the rollout, while the latter remains constant over time [Emmons et al., 2021, Chen et al., 2021]. However, using these conditioning variables during evaluation gives rise to two main issues: (a) the Monte Carlo estimate of return used to compute the cumulative reward-to-go exhibits high variance and (b) the constant average reward-to-go target introduces high bias over time. Based on our analyses of the bias and variance of these approaches in Appendix C, we observe that these issues contribute to decreased performance and stability across runs when evaluating RvS methods.
Although a potential approach to mitigate these issues is to leverage TD learning, such as the Q-Learning Transformer [Yamagata et al., 2022], we restrict our work to RvS methods utilizing behavioral cloning objectives due to the inherent complexity of training value-based methods. To address the aforementioned concerns, we introduce a reward waypoint network denoted as  $W_{\phi}$ , parameterized by  $\phi$ . This network predicts the average and cumulative reward-to-go (ARTG, CRTG) conditioned on the return,  $\omega$ , and current state,  $s_t$ , using offline data  $\mathcal{D}$ . To optimize this network, we minimize the objective shown in Equation 2 using an MSE loss:
$$
\arg \min  _ {\phi} \sum_ {\tau \in \mathcal {D}} \left(\left[ \frac {1}{T - t} \sum_ {t ^ {\prime} = t} ^ {T} \gamma^ {t} r _ {t} \quad \sum_ {t ^ {\prime} = t} ^ {T} \gamma^ {t} r _ {t} \right] ^ {\top} - W _ {\phi} \left(s _ {t}, \omega\right)\right) ^ {2}. \tag {2}
$$
By modeling both ARTG and CRTG, we address the high bias of a constant ARTG target and reduce the variance associated with Monte Carlo estimates for CRTG. The construction of the reward waypoint network is similar in motivation and the prediction task to a baseline network, used to mitigate high variance in methods like REINFORCE [Sutton and Barto, 1999]. However, the distinguishing feature of our reward waypoint network lies in its conditioning on the return, which allows for favorable performance even on suboptimal offline datasets.
5 Waypoint Transformer
We propose the waypoint transformer (WT), a transformer-based offline RL method that leverages the proposed waypoint network  $W_{\phi}$  and a GPT-2 architecture based on multi-head attention [Radford et al., 2019]. The WT policy  $\pi_{\theta}$  is conditioned on past states  $s_{t - k..t}$  and waypoints (either generated goals or rewards)  $\Phi_{t - k..t} = W_{\phi}(s_{t - k..t}, \omega)$  with a context window of size  $k$ , as shown in Figure 3.
Figure 3: Waypoint Transformer architecture, where  $\Phi_t = W_\phi(s_t, \omega)$  represents the output of the goal or reward waypoint network.
Table 1: Normalized scores and training time per task on Gym-MuJoCo, AntMaze, and Kitchen tasks, where bold highlighting indicates SOTA performance, defined by a method's average performance being contained within the interval of the method with highest average.
Environment	TD3 + BC	Onestep RL	CQL	IQL	BC	10% BC	RvS-R/G	DT	QDT	WT (Ours)
halfcheetah-medium-v2	48.3 ± 0.3	48.4 ± 0.1	44.0 ± 5.4	47.4 ± 0.2	42.6	42.5	41.6 ± 0.3	42.4 ± 0.2	42.3 ± 0.4	43.0 ± 0.2
hopper-medium-v2	59.3 ± 4.2	59.6 ± 2.5	58.5 ± 2.1	66.2 ± 5.7	52.9	56.9	60.2 ± 3.0	63.5 ± 5.2	66.5 ± 6.3	63.1 ± 1.4
walker2d-medium-v2	83.7 ± 2.1	81.8 ± 2.2	72.5 ± 0.8	78.3 ± 8.7	75.3	75.0	71.7 ± 1.8	69.2 ± 4.9	67.1 ± 3.2	74.8 ± 1.0
halfcheetah-medium-replay-v2	44.6 ± 0.5	38.1 ± 1.3	45.5 ± 0.5	44.2 ± 1.2	36.6	40.6	38.0 ± 0.7	35.4 ± 1.6	35.6 ± 0.5	39.7 ± 0.3
hopper-medium-replay-v2	60.9 ± 18.8	97.5 ± 0.7	95.0 ± 6.4	94.7 ± 8.6	18.1	75.9	73.5 ± 12.8	43.3 ± 23.9	52.1 ± 20.1	88.9 ± 2.4
walker2d-medium-replay-v2	81.8 ± 5.5	49.5 ± 12.0	77.2 ± 5.5	73.8 ± 7.1	26.0	62.5	60.6 ± 6.7	58.9 ± 7.1	58.2 ± 5.1	67.9 ± 3.4
halfcheetah-medium-expert-v2	90.7 ± 4.3	93.4 ± 1.6	91.6 ± 2.8	86.7 ± 5.3	55.2	92.9	92.2 ± 1.2	84.9 ± 1.6	-	93.2 ± 0.5
hopper-medium-expert-v2	98.0 ± 9.4	103.3 ± 1.9	105.4 ± 6.8	91.5 ± 14.3	52.5	110.9	101.7 ± 16.5	100.6 ± 8.3	-	110.9 ± 0.6
walker2d-medium-expert-v2	110.1 ± 0.5	113.0 ± 0.4	108.8 ± 0.7	109.6 ± 1.0	107.5	109.0	106.0 ± 0.9	89.6 ± 38.4	-	109.6 ± 1.0
gym-avg-v2	75.3 ± 4.9	76.1 ± 2.5	77.6 ± 3.4	76.9 ± 5.8	51.9	74.0	71.7 ± 4.9	65.3 ± 10.1	-	76.8 ± 1.2
antmaze-unmaze-v2	78.6	64.3	74.0	87.5 ± 2.6	54.6	62.8	65.4 ± 4.9	53.6 ± 7.3	-	64.9 ± 6.1
antmaze-unmaze-diverse-v2	71.4	60.7	84.0	62.2 ± 13.8	45.6	50.2	60.9 ± 2.5	42.2 ± 5.4	-	71.5 ± 7.6
antmaze-medium-play-v2	10.6	0.3	61.2	71.2 ± 7.3	0.0	5.4	58.1 ± 12.7	0.0 ± 0.0	-	62.8 ± 5.8
antmaze-medium-diverse-v2	3.0	0.0	53.7	70.0 ± 10.9	0.0	9.8	67.3 ± 8.0	0.0 ± 0.0	-	66.7 ± 3.9
antmaze-large-play-v2	0.2	0.0	15.8	39.6 ± 5.8	0.0	0.0	32.4 ± 10.5	0.0 ± 0.0	-	72.5 ± 2.8
antmaze-large-diverse-v2	0.0	0.0	14.9	47.5 ± 9.5	0.0	6.0	36.9 ± 4.8	0.0 ± 0.0	-	72.0 ± 3.4
antmaze-avg-v2	27.3	20.9	50.6	63.0 ± 8.3	16.7	22.5	53.5 ± 7.2	16.0 ± 2.1	-	68.4 ± 4.9
kitchen-complete-v0	-	-	43.8	62.5	65.0	4.0	50.2 ± 3.6	46.5 ± 3.0	-	49.2 ± 4.6
kitchen-partial-v0	-	-	49.8	46.3	38.0	66.0	51.4 ± 2.6	31.4 ± 19.5	-	63.8 ± 3.5
kitchen-mixed-v0	-	-	51.0	51.0	51.5	40.0	60.3 ± 9.4	25.8 ± 5.0	-	70.9 ± 2.1
kitchen-avg-v0	-	-	48.2	53.3 ± 7.5	51.5	36.7	54.0 ± 5.2	34.6 ± 9.2	-	61.3 ± 3.4
average	-	-	63.7	68.3 ± 6.9	40.1	50.6	62.7 ± 5.7	43.8 ± 7.3	-	71.4 ± 2.8
training time (min)	20	20	80	20	10	10	80	150	-	20
For goal-conditioned and reward-conditioned tasks, we train the goal and reward waypoint network  $W_{\phi}$  respectively on offline dataset  $\mathcal{D}$  independently of the policy. To train the WT policy, we use Algorithm 1 to iteratively optimize its parameters  $\theta$ . During this process, the trained weights  $\phi$  of  $W_{\phi}$  are frozen to ensure the interpretability of the waypoint network's generated goal and reward waypoints. To further simplify the design and improve computational efficiency, the WT is not conditioned on past actions  $a_{t - k..t}$  (i.e., unlike the DT) and we concatenate  $\Phi_t$  with  $s_t$  to produce one token per timestep  $t$  instead of multiple tokens as proposed in Chen et al. [2021].