SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We can apply a similar argument to show that  $\pi_W^*$  (i.e., at optimum) must clone the derived probability when a maximum likelihood objective is applied. Hence, for the optimal intermediate goal-conditioned policy  $\pi_W^*$ , we know it obeys:
$$
\pi_ {W} ^ {*} (a _ {t} = a ^ {(2)} \mid s _ {t} = s ^ {(h)}, \Phi_ {t} = s ^ {(h + 1)}) = \frac {1}{K}
$$
Since  $\pi_G^* (a_t = a^{(2)}\mid s_t = s^{(h)},\omega) = \pi_b(a_t = a^{(2)}\mid s_t = s^{(h)}) = 1 - \lambda$  and we choose  $K < \frac{1}{1 - \lambda}\Rightarrow \frac{1}{K} >1 - \lambda$ , we conclude that:
$$
\pi_ {W} ^ {*} \left(a _ {t} = a ^ {(2)} \mid s _ {t}, \Phi_ {t}\right) > \pi_ {G} ^ {*} \left(a _ {t} = a ^ {(2)} \mid s _ {t}, \omega\right)
$$
This concludes the derivation.
In this section, we provide more details about the experiments, including hyperparameter configuration, sources of reported results for each method, and details of each environment (i.e., version). For all experiments on WT, the proposed method, we run 5 trials with different random seeds and report the mean and standard deviation across them. On AntMaze and Kitchen, we use goal-conditioning, whereas reward-conditioning is used for Gym-MuJoCo. For all experiments on DT, including Gym-MuJoCo, we run 5 trials with random initializations using the default hyperparameters proposed in Chen et al. [2021] and used in the official GitHub repository. We are unable to reproduce some of the results demonstrated in Chen et al. [2021] and reported in succeeding work such as Kostrikov et al. [2021], Emmons et al. [2021].
B.1 Environments and Tasks
AntMaze For AntMaze tasks, we include previously reported results for all methods except RvS-G from Kostrikov et al. [2021]. The results for the RvS-G are from Emmons et al. [2021]. We run experiments for DT (reward-conditioned, as per Chen et al. [2021]) and WT across 5 seeds. For all reported results, including WT, AntMaze v2 is used as opposed to AntMaze v0.
FrankaKitchen On Kitchen, we include available reported results from Kostrikov et al. [2021] for all methods except RvS-G and Emmons et al. [2021] for RvS-G, with results omitted for TD3 + BC and Onestep RL as they are not available in other work or provided by the authors. Similarly to AntMaze, we run experiments for DT and WT across 5 seeds. The target goal configuration for WT is "all" (i.e., where all the tasks are solved), per Emmons et al. [2021]. For all reported results, including WT, Kitchen v0 is used.
Gym-MuJoCo On the evaluated locomotion tasks, we use reported results from Kostrikov et al. [2021] for all methods except RvS-R and Emmons et al. [2021] (RvS-R). We run experiments for DT and WT across 5 seeds. The MuJoCo v2 environments are used for all methods.
B.2 WT Hyperparameters
In Table 3, we show the chosen hyperparameter configuration for WT across all experiments. Consistent with the neural network model in RvS-R/G with 1.1M parameters Emmons et al. [2021], the WT contains 1.1M trainable parameters. For the most part, the chosen hyperparameters align closely with default values in deep learning; for example, we use the ReLU activation function and a learning rate of 0.001 with the Adam optimizer.
In Table 4, we show the chosen hyperparameter configuration for the reward and goal waypoint networks across all experiments. The reward waypoint network always outputs 2 values, the ARTG and CRTG. In general, the goal waypoint network outputs the same dimension as the state since
Table 3: Hyperparameters and configuration details for WT across all experiments.
Hyperparameter	Value
Transformer Layers	2
Transformer Heads	16
Dropout Probability (attn)	0.15
Dropout Probability (resid)	0.15
Dropout Probability (embd)	0.0
Non-Linearity	ReLU
Learning Rate	0.001
Gradient Steps	30,000
Batch Size	1024
it makes  $k$ -step predictions. Depending on the environment, the goal waypoint outputs either a 2-dimensional location for AntMaze or a 30-dimensional state for Kitchen.
Table 4: Hyperparameters and configuration details for goal and reward waypoint networks across all experiments.
Hyperparameter	Value
Number of Layers	3
Dropout Probability	0.0
Non-Linearity	ReLU
Learning Rate	0.001
Gradient Steps	40,000
Batch Size	1024
B.3 Evaluation Return Targets
The target return for the Gym-MuJoCo tasks are specified in Table 5, in the form of normalized scores. These were obtained typically by performing exhaustive grid searches over 4-6 candidate target return values, following prior work [Chen et al., 2021, Emmons et al., 2021]. Typically, we choose the range of the grid search based on the interval close to or higher than the state-of-the-art normalized scores on each of the tasks.
Table 5: Normalized score targets for WT on reward-conditioned tasks in Gym-Mujoco.
Task	Normalized Score Target
hopper-medium-replay-v2	95
hopper-medium-v2	73.3
hopper-medium-expert-v2	125
walker2d-medium-replay-v2	90
walker2d-medium-v2	85
walker2d-medium-expert-v2	122.5
halfcheetah-medium-replay-v2	45
halfcheetah-medium-v2	52.5
halfcheetah-medium-expert-v2	105
C Analysis of Bias and Variance of Reward-Conditioning Variables
We analyze the bias and variance of existing reward-conditioning techniques: a constant average reward-to-go (ARTG) target, as used in Emmons et al. [2021], and a cumulative return target updated with rewards collected during the episode (CRTG), as in Chen et al. [2021]. By analyzing the bias and variance of these techniques, we can determine the potential issues that may explain the performance of methods that condition using these techniques.
Consider the definitions of the true ARTG  $(R_{a})$  and CRTG  $(R_{c})$  below, based on a given trajectory  $\tau$  where the length of the trajectory  $|\tau| = T$ . These definitions are used to train the policy,
$$
R _ {a} (\tau , t) = \frac {1}{T - t} \sum_ {t ^ {\prime} = t} ^ {T} \gamma^ {t} r _ {t} \tag {3}
$$
$$
R _ {c} (\tau , t) = \sum_ {t ^ {\prime} = t} ^ {T} \gamma^ {t} r _ {t} \tag {4}
$$
At evaluation time, it is impossible to calculate  $r_{t'}$  for any  $t' \geq t$ . As a result, we provide ARTG and CRTG targets,  $\theta_a$  and  $\theta_c$  respectively. At evaluation time, the values of the ARTG and CRTG are estimated and used as follows in Chen et al. [2021] and Emmons et al. [2021]:
$$
\hat {R} _ {a} (\tau , t) = \theta_ {a} \tag {5}
$$
$$
\hat {R} _ {c} (\tau , t) = \theta_ {c} - \sum_ {t ^ {\prime} = 1} ^ {t} \gamma^ {t} r _ {t} \tag {6}
$$
Ideally, the errors of the estimated  $\hat{R}_a$  and  $\hat{R}_c$  are minimal so as to accurately approximate the true average or cumulative reward-to-go respectively, but that is often infeasible. To characterize the error of each of the evaluation estimates of ARTG and CRTG, consider the decomposition of the error for a particular reward conditioning variable  $R$  and estimated evaluation  $\hat{R}$  presented in Theorem C.1.
Theorem C.1. The general bias-variance decomposition of the expected squared error between a true  $R(\tau, t)$  (e.g., Equations 1 or 2) and an estimator  $\hat{R}(\tau, t)$  (e.g., Equations 3 or 4) under the trajectory distribution induced by an arbitrary policy  $\pi$  and unknown transition dynamics  $p(s_{t+1} \mid s_t, a_t)$  is given by:
$$
\mathbb {E} _ {\tau} \left[ \left(R (\tau , t) - \hat {R} (\tau , t)\right) ^ {2} \right] = \mathbb {E} \left[ R (\tau , t) - \hat {R} (\tau , t) \right] ^ {2} + \operatorname {V a r} [ \hat {R} (\tau , t) - R (\tau , t) ]
$$
Proof. Similarly to the derivation of the standard bias-variance tradeoff, we expand terms and separate into multiple expectations using the linearity of expectation. We leverage the definition of the covariance,  $\mathrm{Cov}(X,Y) = \mathbb{E}[XY] - \mathbb{E}[X]\mathbb{E}[Y]$ , and variance,  $\mathrm{Var}[X] = E[X^2] - E[X]^2$ , in several steps.
$$
\mathbb {E} _ {\tau} \left[ \left(R (\tau , t) - \hat {R} (\tau , t)\right) ^ {2} \right] = \mathbb {E} \left[ R (\tau , t) ^ {2} \right] + \mathbb {E} \left[ \hat {R} (\tau , t) ^ {2} \right] - 2 \mathbb {E} \left[ \hat {R} (\tau , t) R (\tau , t) \right]
$$
We can simplify the first term using the definition of the variance and the third term using the definition of the covariance.
$$
\begin{array}{l} \mathbb {E} _ {\tau} [ (R (\tau , t) - \hat {R} (\tau , t)) ^ {2} ] = \operatorname {V a r} [ R (\tau , t) ] + \mathbb {E} [ R (\tau , t) ] ^ {2} + \mathbb {E} [ \hat {R} (\tau , t) ^ {2} ] - 2 \mathbb {E} [ \hat {R} (\tau , t) R (\tau , t) ] \\ = \operatorname {V a r} [ R (\tau , t) ] + \mathbb {E} [ R (\tau , t) ] ^ {2} + \mathbb {E} [ \hat {R} (\tau , t) ^ {2} ] - 2 (\operatorname {C o v} (\hat {R} (\tau , t), R (\tau , t)) + \\ \mathbb {E} [ R (\tau , t) ] \cdot \mathbb {E} [ \hat {R} (\tau , t) ]) \\ \end{array}
$$
Similarly, we simplify the  $\mathbb{E}[\hat{R} (\tau ,t)^2 ]$  term using the definition of the variance and collect terms.
$$
\begin{array}{l} \begin{array}{c}\mathbb{E}_{\tau}[(R(\tau ,t) - \hat{R} (\tau ,t))^{2}] = (\mathbb{E}[R(\tau ,t)] - \mathbb{E}[\hat{R} (\tau ,t)])^{2} + \operatorname {Var}[\hat{R} (\tau ,t)] - 2\operatorname {Cov}(\hat{R} (\tau ,t),R(\tau ,t)) + \\ \operatorname {Var}[R(\tau ,t)] \end{array} \\ = \begin{array}{c} \mathbb {E} [ R (\tau , t) - \hat {R} (\tau , t) ] ^ {2} + \operatorname {V a r} [ \hat {R} (\tau , t) ] - 2 \operatorname {C o v} (\hat {R} (\tau , t), R (\tau , t)) + \\ \operatorname {V a r} [ R (\tau , t) ] \end{array} \\ \end{array}
$$
Equivalently, since  $\operatorname{Var}[X - Y] = \operatorname{Var}[X] + \operatorname{Var}[Y] - 2\operatorname{Cov}(X,Y)$ , we can rewrite the result as follows.
$$
\mathbb {E} _ {\tau} \left[ \left(R (\tau , t) - \hat {R} (\tau , t)\right) ^ {2} \right] = \mathbb {E} \left[ R (\tau , t) - \hat {R} (\tau , t) \right] ^ {2} + \operatorname {V a r} \left[ \hat {R} (\tau , t) - R (\tau , t) \right]
$$
Figure 7: True average reward-to-go as a function of timestep  $t$  for 200 rollouts of a transformer policy on hopper-medium-replay-v2 compared to the constant ARTG target  $\theta_{a}$  (dotted black line), for left: successful rollouts and right: unsuccessful rollouts.
This completes the derivation of the general bias-variance decomposition.
Analysis of ARTG Consider the decomposition of the error per the bias-variance tradeoff of the ARTG. Trivially, the variance of the estimate in Equation 3 is zero and it is independent of  $R$  because it is a constant value, and as a result, the expected error is composed entirely of the bias and irreducible variance.
$$
\begin{array}{l} E _ {\tau} \left[ \left(R _ {a} (\tau , t) - \hat {R} _ {a} (\tau , t)\right) ^ {2} \right] = E _ {\tau} \left[ R _ {a} (\tau , t) - \hat {R} _ {a} (\tau , t) \right] ^ {2} + \operatorname {V a r} \left[ R _ {a} (\tau , t) \right] \\ = E _ {\tau} \left[ R _ {a} (\tau , t) - \theta_ {a} \right] ^ {2} + \operatorname {V a r} \left[ R _ {a} (\tau , t) \right] \\ \end{array}
$$
Based on the terms that we are able to minimize, we can derive through calculus that the squared error and bias of a constant estimator are minimized by the mean, i.e., when  $\hat{\theta}_a = \mathbb{E}_{\tau}[R_a(\tau ,t)]$ . However, in the offline RL setting, we cannot easily determine this value for an arbitrary trained policy  $\pi$ . As a result, the bias of the technique during evaluation can lead to high error in estimating the true ARTG, which may consequently cause reduced performance during evaluation of the policy.
Empirically, we demonstrate that the high bias of this technique may lead to instability in achieved return across rollouts on hopper-medium-replay-v2. Specifically, suppose that a rollout is classified as unsuccessful if it terminates before the time limit  $T = 1000$ , and analogously, a successful rollout reaches  $t = T$  without termination. Based on these distinctions, we display the true ARTG across 200 successful and unsuccessful rollouts of a trained transformer policy in Figure 7.
Clearly, for all successful rollouts, the true ARTG matches the constant target closely across most  $t \in [1, T)$ , with the exception of  $t \approx T$ . However, across most failed rollouts, the bias of the constant target for small  $t \approx 0$  (i.e., when the policy is taking its first few actions) is relatively high. As  $t$  grows, the bias seems to spike upwards significantly and the episode terminates shortly thereafter, indicating an unsuccessful rollout.
Hence, it is evident that when the true ARTG closely follows the prescribed constant ARTG target, the policy consistently achieves state-of-the-art performance. However, whenever the target ARTG underestimates or overestimates the true ARTG by a margin of greater than 0.4, the rollout tends to be unsuccessful, often achieving less than half the return compared to successful rollouts.
To that end, the reward waypoint network uses a neural network formulation to estimate  $R_{a}$  (i.e., without using a constant estimator  $\theta_{a}$ ) based on the state  $s_t$  and the target return  $\omega$ . By training the neural network to provide a less biased estimate of the ARTG, we show that we can achieve substantial performance improvements over a constant estimate of ARTG on the same task. As shown in Table 6, RvS-R and a constant ARTG both exhibit significantly lower average performance and greater variability compared to WT (with a reward waypoint network).
Table 6: Normalized evaluation scores for different policies and ARTG estimation techniques on the hopper-medium-replay-v2 task.
Technique	Normalized Score
Transformer (constant ARTG)	66.5 ± 15.6
RvS-R (constant ARTG)	73.5 ± 12.8
WT (waypoint network)	88.9 ± 2.4
Analysis of CRTG Consider a similar decomposition of the error using Theorem C.1 of the CRTG. Though neither the bias nor the variance of  $\hat{R}_c$  are necessarily zero, it is important to note that the variance term can be large even if the bias is zero, i.e., if  $\theta_c = \mathbb{E}[\sum_{t'=1}^T \gamma^t r_t]$ .
Corresponding to a best-case scenario where  $\hat{R}_c$  is unbiased, we consider a bound of the expected error constructed using only the variance term. For simplicity, we assume that we incur the worst-case variance term with minimal covariance between  $R_c$  and  $\hat{R}_c$ . The resulting dominant quantities are the variance of  $\hat{R}_c$  and the irreducible variance of  $R_c$ .
$$
\begin{array}{l} E _ {\tau} \left[ \left(R _ {c} (\tau , t) - \hat {R} (\tau , t)\right) ^ {2} \right] = \mathbb {E} \left[ R (\tau , t) - \hat {R} (\tau , t) \right] ^ {2} + \operatorname {V a r} \left[ \hat {R} (\tau , t) - R (\tau , t) \right] \\ \geq \operatorname {V a r} [ \hat {R} (\tau , t) - R (\tau , t) ] \\ \approx \operatorname {V a r} \left[ \sum_ {t ^ {\prime} = 1} ^ {t} \gamma^ {t ^ {\prime}} r _ {t ^ {\prime}} \right] + \operatorname {V a r} [ R (\tau , t) ] \\ \end{array}
$$
Prior work has shown that Monte Carlo estimates of return exhibit high variance, motivating techniques such as using baseline networks for REINFORCE,  $k$ -step returns or TD( $\lambda$ ) for TD-learning methods, etc., which aim to reduce variance at the cost of potentially incurring bias Sutton and Barto [1999]. In that vein, we demonstrate that our reward waypoint network predicts a similar quantity as a baseline network and inherits many of its desirable properties.
Specifically, while the baseline network predicts  $\sum_{t'=t}^{T} \gamma^{t} r_{t}$  given  $s_{t}$ , we additionally condition our reward waypoint network's prediction on  $\omega$ . In the offline RL setting where datasets are often a mixture of suboptimal and optimal trajectories, we require that the waypoint network can differentiate between such trajectories. Hence, by conditioning on the overall reward  $\omega$ , the reward waypo high and low return-achieving trajectories.
Figure 8: Variance of estimated CRTG (Equation 4) and predicted CRTG by the reward waypoint network on the hopper-medium-replay-v2.
Empirically, we demonstrate that the predicted CRTG from our reward waypoint network exhibits lesser variance on hopper-medium-replay-v2 than the estimated CRTG (i.e., as in Equation 4). As shown in Figure 8, the variance of both techniques are relatively similar until  $t \approx 400$ , after which the variance of the estimated CRTG appears to grow superlinearly as a function of  $t$ . In the worst case, at  $t \approx T$ , the variance is nearly 5x larger than for the predicted CRTG.
D.1 Analysis of Stitching Region Behavior
To add on to the analysis of the goal waypoint network presented in the main text, we analyze the "failure" regions of transformer policies with and without a goal waypoint network. That is, by determining the final locations of the agent, we can examine where the agent ended up instead of the target location. Similar to the analysis in Section 6.3, this analysis can inform the stitching capability of our methods.
Figure 9: End locations for antmaze-large-play-v2 during 100 rollouts of left: WT and right: global goal-conditioned transformer policy.
Based on Figure 9, it is clear that the WT does not get "stuck" (e.g., after taking the wrong turn) as often as the policy conditioned on global-goals. Moreover, the number of ants ending up near the beginning portions of the maze (i.e., the bottom left) is significantly smaller for WT, which contributes to its doubled success rate. We believe these are primarily attributable to the guidance provided by the goal waypoint network through a consistent set of intermediate goals to reach the target location at evaluation time.
Interestingly, we observe that WT displays an increased rate of failure around the final turn relative to other regions in the maze. As there is a relative lack of density in other failure regions closer to the beginning of the maze, we hypothesize that some rollouts may suffer from the ant getting "stuck" at challenging critical points in the maze, as defined in Kumar et al. [2022]. This indicates an interesting direction of exploration for future work and a technique to combat this could result in policies with nearly  $100\%$  success rate in completing antmaze-large-play-v2.
D.2 Target Reward Interpolation
Unlike traditional value-based approaches, RvS methods such as DT, RvS, and WT present a simple methodology to condition the policy to achieve a particular target return. In theory, this allows RvS to achieve performance corresponding to the desired performance level, provided that it has non-zero coverage within the offline training dataset.
To examine RvS's capability in this regard, we replicate an analysis performed in Emmons et al. [2021] to examine whether RvS can interpolate between modes in the data to achieve a particular target return. In this case, we compare WT to RvS-R on the walker2d-medium-expert-v2 task, in which the data is composed of two modes of policy performance (i.e., medium and expert). The mode corresponding to the "medium" data is centered at a normalized score of 80, whereas "expert" performance is located at a normalized score of 110.
Based on Figure 10, WT shows a reasonably improved ability to interpolate in some regions than RvS-R. Where RvS-R displays a failure to interpolate from normalized targets of 30-60 and WT does not, both tend to be unable to interpolate between 70-100 (i.e., between the modes of the dataset).
Figure 10: Achieved normalized score versus the target normalized score for walker2d-medium-expert-v2 task on WT and RvS-R.
Although the reasons for the failure in interpolation are unclear across two RvS methods, it is a worthwhile analysis for future work. We hypothesize that techniques such as oversampling may mitigate this issue as this may simply be linked to low frequency in the data distribution.
D.3 Comparisons to Manual Waypoint Selection
We compare the performance of the proposed goal waypoint network with a finite set of manual waypoints, hand-selected based on prior oracular knowledge about the critical points within the maze for achieving success (i.e., turns, midpoints). Based on the selected manual waypoints, shown in Figure 11, we use a simple algorithm to provide intermediate targets  $\Phi_t$  based on a distance-based sorting approach, shown in Algorithm 2.
Figure 11: Manually selected waypoints (blue pluses) for antmaze-large-play-v2, the chosen task to evaluate the proposed approach. As before, the start location is marked with a maroon dot, and the target location is marked wit a gold star.
Algorithm 2 Manual waypoint selection with  $W_{m}$  and  $s_t$  using  $L_{2}$  distance and a given global goal  $\omega$ .
$W_{c}\gets \{w_{m}:||w_{m} - \omega ||_{2}\leq ||s_{t} - \omega ||_{2}\}$  {consider waypoints that brings agent closer to  $\omega$  } return arg  $\min_{w_c\in W_c}\| w_c - s_t\| _2$
With all configuration and hyperparameters identical to WT, we compare the performance of a global goal-conditioned policy, WT with manual waypoints, and WT with the goal waypoint network on antmaze-large-play-v2 in Table 7.
The results demonstrate that WT clearly outperforms manual waypoint selection in succeeding in the AntMaze Large environment. However, while comparing a global-goal conditioned policy and a policy conditioned on manual waypoints, it is clear that the latter improves upon average performance and variability across initialization seeds. We believe that this illustrates that (a) waypoints, whether manual or generated, tend to improve performance of the policy and (b) finer-grained waypoints provide more valuable information for the policy to succeed more.
Table 7: Normalized evaluation scores for different waypoint selection techniques on the antmaze-large-play-v2 task.
Technique	Normalized Score
No Waypoints	33.0 ± 10.3
Manual Waypoints	44.5 ± 2.8
Waypoint Network	72.5 ± 2.8
We believe that this provides further verification and justification for both the generation of intermediate targets and the procedure of generation through a goal waypoint network that performs  $k$ -step prediction.
D.4 Delayed Rewards
An important case to consider for reward-conditioned tasks is when the rewards are delayed (often, provided at the end of the trajectory). By providing nearly no intermediate targets, it is often challenging to complete these tasks. To verify that our design choices to construct modeled intermediate
targets does not reduce performance on configurations in which such information is unavailable, we evaluate WT on 3 MuJoCo tasks where all rewards are provided at the final timestep. Moreover, we ablate the choice of removing actions from WT, which may have an effect on the performance considering that intermediate rewards are not provided.
To compare our performance with other methods, we report results for CQL, DT, and QDT from Yamagata et al. [2022] for all non-expert tasks and from Chen et al. [2021] from hopper-medium-expert-v2 in Table 8. The results demonstrate that our method is relatively comparable to DT in performance and that removing actions does not significantly affect performance in these cases. All BC methods significantly outperform CQL across the chosen delayed reward tasks.
Table 8: Normalized scores across 5 seeds on delayed reward MuJoCo v2 tasks, where all rewards are provided at the final timestep, on WT with action conditioning  $(a_{t})$  and without action conditioning.
Environment	CQL	DT	QDT	WT (at)	WT
hopper-medium-expert-v2	9.0	107.3 ± 3.5	-	103.9 ± 3.0	104.2 ± 2.4
halfcheetah-medium-replay-v2	7.8 ± 6.9	33.0 ± 4.8	32.8 ± 7.8	30.6 ± 2.6	30.9 ± 2.8
walker2d-medium-v2	0.0 ± 0.4	69.9 ± 2.0	63.7 ± 6.4	71.5 ± 0.9	70.8 ± 1.7
E Baseline Reproduction
We reproduce the TD3+BC results reported in Fujimoto and Gu [2021] and subsequently in Kostrikov et al. [2021], Emmons et al. [2021], which are not reproduced by Liu and Abbeel [2023]. For 8 MuJoCo tasks chosen from the same evaluation tasks as WT, we show the normalized score of TD3+BC and the reported mean score as a function of the number of iterations of training. For all of the tasks, it is evident that our reproduction aligns with the original reported result, whereas those in Liu and Abbeel [2023] are consistently higher. We hypothesize that these findings are similar in nature to our inability to reproduce results for DT in 1.