SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Many prior work [48, 18] reports their performance by taking the final checkpoint's evaluation results and averaging them over multiple seeds. However, in some cases we observe that this evaluation protocol can lead to large oscillations in the policy performance when gradient steps (or training epochs) are slightly perturbed or the random seeds are changed (as in Figure A1). As different methods may use various number of training steps, seeds and evaluation trajectories, it hampers drawing robust conclusions by comparing these results. To faithfully capture the performance trend, which is crucial for ablation analysis, for each seed we take the best evaluation result among the last five checkpoints as its performance, where each evaluation result is the average return over  $N$  evaluation trajectories (e.g.  $N = 10, 50, 100$ ) achieved by a given checkpoint, and then we calculate the mean performance and standard deviation over five seeds. Empirically we find this evaluation protocol yields results that are less sensitive to the above factors and usually have smaller error bounds, facilitating the comparisons in ablation studies. Meanwhile, we observe that in some cases this evaluation protocol might lead to results that are higher than those reported in previous papers, especially when the policy has more drastic performance fluctuations during training. We re-run most baselines using the same evaluation protocol for fair comparisons.
To investigate the impact of different aggregation metrics, we follow [1] to compare performance aggregation using mean, median, IQM and optimality gap on the AntMaze tasks. In Figure A2, we observe that GCPC consistently outperforms compared baselines across all metrics.
$^{1}$ https://github.com/jhejna/dwsl/
Figure A1: Training curves of different methods with various default gradient steps
Figure A2: Different aggregation metrics with  $95\%$  confidence intervals on AntMaze tasks.
Normalized Score
Apart from model-free policies, it also remains unclear whether or how the results of planning-based methods [2, 27] would be affected by the choice of evaluation protocol. We leave adopting more reliable and robust evaluation approaches for comprehensive analysis to the future work.
A.3 Hyperparameter
We list hyperparameters for BC and RvS-G/R replication in Table A1 and GCPC implementation in Table A2. For RvS-G/R and PolicyNet in GCPC, we use a two-layer feedforward MLP as the policy network, taking the current state and goal (state or return-to-go) as input, the only difference is that PolicyNet takes the bottleneck as an additional input. All experiments are performed on a single Nvidia RTX A5000.
Table A1: Replicated BC and RvS-G/R Hyperparameter
Hyperparameters	Value
Hidden Layers	2
Embedding Dim	1024
Nonlinearity	ReLU
Epochs	80 AntMaze, Kitchen
100 Gym Locomotion
Batch Size	1024 AntMaze, Gym Locomotion
256 Kitchen
Dropout	0
Optimizer	Adam
Learning Rate	1e-3
A.4 Explicit future
Existing methods, like TT and DD, utilize future information by sampling future states (or transitions) and planning on the explicit future. TT learns a world model that autoregressively generates states, actions, rewards and return-to-go estimates, then uses beam search to select future trajectories
Table A2: GCPC Hyperparameter
Hyperparameters	Values
Optimizer	Adam
TrajNet(Bidirectional Transformer)	Encoder Layers	2
Decoder Layers	1
Attention Heads	4
Embedding Dim	256
Nonlinearity	GELU
Slot Tokens	4
(History Window k, Future Window p)	(10, 70) AntMaze
(5, 30) Kitchen
(5, 20) Locomotion
Epochs	60 AntMaze, Locomotion
20 Kitchen
Batch Size	1024
Dropout	0.1
Learning Rate	1e-4
PolicyNet(Multi-Layer Perceptron)	Hidden Layers	2
Embedding Dim	1024
Nonlinearity	ReLU
Epochs	80 AntMaze, Kitchen
100 Gym Locomotion
Batch Size	1024 AntMaze, Gym Locomotion
256 Kitchen
Dropout	0
Learning Rate	1e-3 AntMaze, Gym Locomotion
1e-4 Kitchen
with highest cumulative rewards. DD leverages the diffusion process with classifier-free guidance to generate future states that satisfy reward maximization or other constraints, and uses inverse dynamics to extract actions between future states.
In GCPC, we compress the goal-conditioned future into latent representations and use them to support policy learning. In Section 4.4, we study the effect of explicitly leveraging the decoded future states. We adapt our framework as follows: During the policy learning and evaluation, we first freeze transformer encoder and decoder trained in the first stage. Instead of directly passing the encoded bottlenecks to the PolicyNet, we leverage the pre-trained transformer decoder to decode the bottleneck into goal-conditioned futures. The new PolicyNet is modified to take the observed states, the goal, and the concatenation of the decoded future states. The bottleneck is discarded and not used by the modified PolicyNet.
A.5 Future window length
We further investigate the effect of future window length on policy performance, which decides how far the bottleneck can look into the future during the first stage pretraining. In Table A3, we show that predicting further future would generally bring bigger improvements in AntMaze-Large, emphasizing that planning into long-term future can largely benefit policy performance. We also observe that the improvements brought by enlarging future window are specific to both datasets and the range of future window length. The improvements become marginal once the future window expands beyond a certain size.
A.6 Goal Sampling and Specification
To perform goal-conditioned imitation learning, we follow [15] to sample an outcome from the future trajectory  $\tau_{t:H}$  as the goal. When the goal is in the form of a target state (e.g. in AntMaze and Kitchen tasks), we randomly sample a reachable future state from  $\{s_i\}_{i = t + 1}^H$  and keeps the task-related subspace as  $g$ . Specifically, for AntMaze we take the ant's  $(x,y)$  location as the goal, which are the first two dimensions of the state space; for Kitchen's goal, we keep the dimensions
Table A3: Future window length. In AntMaze-Large, the farther into the future can the bottleneck see during trajectory representation pre-training, the more helpful they are for policy learning.
Future window length p	20	40	70
Large-Diverse	62.4 ± 5.6	78.2 ± 4.8	80.6 ± 3.9
Large-Play	49.6 ± 6.7	76.6 ± 5.0	78.2 ± 3.2
Average	56.0	77.4	79.4
that indicate the position of the target object in all 7 subtasks and zero out other dimensions (e.g. dimensions that indicate the robot proprioceptive state) in the 30-dimensional state space<sup>2</sup>. When the goal is represented as the expected return-to-go (e.g. in Gym Locomotion tasks), we use the average return achieved over some number of timesteps into the future as  $g$  (i.e.  $g = \frac{1}{H - t + 1}\sum_{i = t}^{H}r_{i}$ , where we follow [15] to set  $H$  to the constant maximum episode length).
B Pseudocode of GCPC
Algorithm 1 Goal-conditioned Predictive Coding (GCPC) for RvS
1: Input: Dataset of trajectories  $\mathcal{D} = \{\tau\}$ , History window length  $k$ , Future window length  $p$ , Masking ratio  $r$ , Bi-directional transformer encoder  $h_{\mathrm{enc}}$ , Bi-directional transformer decoder  $h_{\mathrm{dec}}$ , Learnable slot tokens  $Q$ , Policy network  $\pi$
2: while First Stage do
3: Sample a trajectory and a timestep:  $\tau = \{(s_i, a_i)\}_{i=1}^H \sim \mathcal{D}, t \sim [k, H]$
4: Sample a goal  $g$ , obtain observed states and full states:  $s_h = s_{t - k + 1:t}$ ,  $s_w = s_{t - k + 1:t + p}$
5: Randomly mask observed states with masking ratio  $r$ :  $s_h^m = \text{Masked}(s_h, r)$
6: Encode masked observed states, goal and learnable slot tokens:  $B = h_{\mathrm{enc}}(Q, s_h^m, g)$
7: Reconstruct observed and future states with the bottleneck  $B$ :  $\hat{s}_w = h_{\mathrm{dec}}(B)$
8: Compute loss:  $\mathcal{L} = \mathrm{MSE}(\hat{s}_w,s_w)$
9: Update parameters in  $h_{\mathrm{enc}}$  and  $h_{\mathrm{dec}}$
10: end while
11: Freeze parameters in  $h_{\mathrm{enc}}$
12: while Second Stage do
13: Sample a trajectory and a timestep:  $\tau = \{(s_i, a_i)\}_{i=1}^H \sim \mathcal{D}, t \sim [k, H]$
14: Sample a goal  $g$ , obtain observed states:  $s_h = s_{t - k + 1:t}$
15: Encode observed states, goal and learnable slot tokens:  $B_{\mathrm{obs}} = h_{\mathrm{enc}}(Q,s_h,g)$
16: Train the policy network  $\pi$  with the bottleneck  $B_{\mathrm{obs}}$ :  $\hat{a}_t = \pi (B_{\mathrm{obs}}, s_t, g)$
17: Compute loss:  $\mathcal{L} = \mathrm{MSE}(\hat{a}_t, a_t)$
18: Update parameters in  $\pi$
19: end while
20: return Policy network  $\pi$
C Pre-training mask ratio
In the first stage of GCPC, we randomly mask history states with a masking ratio  $r$ . We compare using fixed masking ratios (20%, 40%, 80%) and dynamic masking ratio "D" uniformly sampled from 0%, 20%, 40%, 60% and 80%. Figure A3 shows the influence of different masking ratios on AntMaze and FrankaKitchen. We found that low masking ratios (e.g. 20%) would usually perform worse except on the Kitchen-Mixed dataset, and a high masking ratio (80%) turns out to be more effective, which is similar to the observation in MAE [24]. However, the optimal masking ratio is not consistent across environments. In order to achieve the best overall performance, we adopt dynamic masking ratio "D" in our implementation.
In the Kitchen experiments, we found it important to zero-out redundant dimensions, such as robot proprioceptive dimensions, when specifying the task goal for goal (state)-conditioned methods. We use this specification for all goal (state)-conditioned methods.
Figure A3: Mask ratio ablation on AntMaze and FrankaKitchen. “D” refers to dynamic mask ratio randomly drawn from  $0\%$ ,  $20\%$ ,  $40\%$ ,  $60\%$  and  $80\%$
D Additional Experiments on FrankaKitchen
In Table A4, we present how different pretraining objectives affect policy performance on the FrankaKitchen benchmark. MAE-F and MAE-RC are the only valid objectives in zero-shot setting. In two-stage setting, we observe that action inputs would generally cause larger variance and harm the performance of two predictive coding objectives. Similar to our observations on AntMaze, predictive coding objectives can consistently generate powerful trajectory representations, which yield performance gains on Kitchen tasks.
Table A4: Objective Ablation on Kitchen-Partial. We report mean and standard error over 5 seeds.
Kitchen-Partial	AE-H	MAE-H	MAE-F	MAE-RC	MAE-ALL
Zero-shot	-	0	4.4 ± 9.8	4.4 ± 6.5	0
Two-stage (w/ actions)	76.9 ± 23.9	73.6 ± 16.6	86.4 ± 12.5	80.4 ± 18.0	66.8 ± 14.5
Two-stage (w/o actions)	66.4 ± 5.9	72.7 ± 7.0	92.5 ± 4.9	90.2 ± 6.6	68.0 ± 4.9
E Visualizing the Latent Future
In this section, we visualize the latent future by decoding the bottlenecks back into the state space using the pre-trained decoder. We visualize the currently observed state and the decoded future states inside the maze environment. Figure A4 includes two examples extracted from successful evaluation rollouts. The blue part on the left shows the actual position of the agent at a certain timestep  $s_t$  during the evaluation and the green part on the right is the latent future predicted based on history states  $s_{\leq t}$  and the goal  $g$ . We find in the latent future, the agent gradually move towards the goal, which demonstrates that bottlenecks can encode meaningful future trajectories that provide the correct waypoints as subgoals and bring the agent from the current location to the final goal.
Figure A5 illustrates how latent future vary depending on different goal conditions for the same current state. By specifying different goals, the latent future would change correspondingly and point out the right direction for the agent. These visualizations suggest that the bottlenecks have the ability to provide the correct intermediate waypoints based on history states and the goal, which is crucial for long-horizon goal-reaching tasks. In the second example of Figure A5, multiple paths connect the agent and the goal, but bottlenecks encode the shortest one for the agent, which also demonstrates the bottleneck's ability to stitch sub-optimal trajectories and compose the optimal path.
In Section 4.3, we found that removing the goal from TrajNet input would seriously harm the performance on large AntMaze. In Figure A6, we visualize the latent future without goal conditioning. In both examples, bottlenecks encode latent futures that either move in the opposite direction to the goal or enter a dead end and cause tasks to fail. It again demonstrates that the quality of the encoded
Goal 1 Location
Goal 2 Location
Agent Location
Latent Future for Timestep t1 (w/ Goal 1)
Figure A4: Two latent future examples from successful rollouts. Latent future shows the right path to the goal.
Latent Future for Timestep t2 (w/ Goal 2)
Figure A5: Latent future conditioned on different goals in large AntMaze.
future would significantly affect how policy performs, and providing goal information can prevent misleading future.
Goal 3 Location
Goal 4 Location
Timestep t3
Latent Future for Timestep t3 (w/o Goal 3)
Evaluation
Timestep t4
Latent Future for Timestep t4 (w/o Goal 4)
Figure A6: Failure cases after removing goal from TrajNet input. Latent future without goal conditioning would mislead the agent.
F Slot tokens ablation
Table A5 shows how changing the number of slot tokens would affect policy performance on large AntMaze. We observed that using different numbers of slot tokens would not have big impacts on the performance. In our implementation, we uniformly use 4 slot tokens for both trajectory representation learning and policy learning.
Table A5: The number of slot tokens. We report mean and standard error over 3 seeds.
# Slot Tokens	1	2	4
Large-Diverse	80.6 ± 3.5	77.2 ± 4.9	80.6 ± 3.9
Large-Play	81.2 ± 4.4	78.8 ± 4.7	78.2 ± 3.2