SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To answer the questions above, we conduct extensive experiments on three domains from D4RL offline benchmark suite [16]: AntMaze, FrankaKitchen and Gym Locomotion. AntMaze is a class of long-horizon navigation tasks, featuring partial observability, sparse reward and datasets that consist primarily of suboptimal trajectories. In this domain, an 8-DoF Ant robot needs to "stitch" parts of subtrajectories and navigates to a particular goal location in partially observed mazes. In addition to three mazes from the original D4RL, we also include a larger maze (AntMaze-Ultra) proposed by [28]. Both large and ultra setup in AntMaze poses significant challenges due to the complex maze
Figure 3: D4RL environments used in evaluation. From left to right: AntMaze, FrankaKitchen, Walker 2D, Hopper, and Halfcheetah.
layout and long navigation horizon. FrankaKitchen is a long-horizon manipulation task, in which a 9-DoF Franka robot arm is required to perform 4 subtasks in a simulated kitchen environment (e.g. open the microwave, turn on the light). In Gym Locomotion, we evaluate our approach on three continuous control tasks: Walker 2D, Hopper and Halfcheetah. Following [15], we refer to the "goal" of AntMaze and Kitchen as the target state configuration, and that of Gym as the average return-to-go.
Experimental details. For all benchmarks, we use a two-layer transformer encoder and a one-layer transformer decoder as the TrajNet, and a two-layer MLP as the PolicyNet (see detailed hyperparameters in A.3). During policy learning, we use the pre-trained TrajNet that achieves the lowest validation reconstruction loss to generate the bottleneck. For each evaluation run, we follow [29] to take the success rate over 100 evaluation trajectories for AntMaze tasks. For Kitchen and Gym Locomotion, we average returns over 50 and 10 evaluation trajectories, respectively. To determine the performance with a given random seed, we take the best evaluation result among the last five checkpoints (see discussion in A.2). For performance aggregation, we report the mean performance and standard deviation averaged over five seeds for each experiment.
4.2 Impact of Trajectory Representation Pretraining Objectives
In the two-stage framework, the pre-trained TrajNet generates trajectory representations in the form of the bottleneck, which is taken as an input to PolicyNet. To study the impact of trajectory representation learning objectives on the resulting policy performance, we implement five different sequence modeling objectives (as in Table 1) by varying masking patterns in the first stage pretraining.
Table 2 compares the performance of policies with different settings and pretraining objectives<sup>1</sup>. We note that MAE-F is the only effective masking pattern to perform zero-shot inference. After decoupling representation learning and policy learning, the MLP policy consistently outperforms the zero-shot transformer policy. This suggests good objectives for two stages could be different – by decoupling we can get the best of both worlds. Also, we observe that removing action sequences from trajectory representation pretraining yields performance gains in this task, whereas previous single-stage Transformer policy (e.g. [10]) usually requires action inputs to function, which further demonstrates the flexibility of our decoupled framework. With properly designed objectives, sequence modeling can generate powerful trajectory representations that facilitate the acquisition of performant policies.
Figure 4: Policy learning curves with different pretraining objectives on Kitchen-Partial.
In both Table 2 and Figure 4, we observe that both MAE-F and MAE-RC outperform the other objectives in both AntMaze-Large and Kitchen-Partial, confirming the importance of the predictive coding objective for trajectory representation learning.
Table 2: Comparison of trajectory representation pretraining objectives. We evaluate five different objectives under three settings on large AntMaze environment: (1) Zero-shot: When actions are considered as part of the masked tokens, the pretrained TrajNet can be directly utilized as the policy; (2) Two-stage (w/o actions): Two-stage framework employs an MLP as PolicyNet, with state-only trajectories as the input to TrajNet. (3) Two-stage (w/ actions): Two-stage framework employs an MLP as PolicyNet, with state-action trajectories as the input to TrajNet.
Large-Play	AE-H	MAE-H	MAE-F	MAE-RC	MAE-ALL
Zero-shot	-	0	21.2 ± 17.1	0	0
Two-stage (w/ actions)	12.6 ± 4.7	6.4 ± 3.8	57.2 ± 5.5	62.0 ± 10.7	11.6 ± 6.3
Two-stage (w/o actions)	30.2 ± 6.6	36.6 ± 12.6	76.2 ± 4.0	78.2 ± 3.2	32.0 ± 13.2
4.3 The Role of Goal Conditioning in Trajectory Representation Pretraining
We also investigate whether goal conditioning (i.e. the goal input) in TrajNet is necessary or beneficial for learning trajectory representations. In Table 3, we observe that the objectives without performing
For pre-training on AntMaze, when actions (i.e. the torque applied on joints) are considered, the full state space is reconstructed; when state-only trajectories are used, we reconstruct only the locations of the agent.
Figure 5: Reconstruction loss on the validation set during the trajectory representation learning stage, when MAE-F objective is used. "GC" refers to goal conditioning in TrajNet.
future prediction is insensitive to the goal conditioning. For example, the results of AE-H and MAE-H pretraining objectives align with the intuition that merely summarizing the history should not require the goal information. However, goal conditioning is crucial for predictive coding objectives (e.g. MAE-F and MAE-RC). Removing the goal from TrajNet might result in aimless prediction and misleading future representations, which thus harm the policy. Figure 5 shows the curves of validation loss during pretraining with MAE-F objective. We observe that goal conditioning reduces the prediction error and helps the bottleneck properly encode the expected long-term future, which would benefit the subsequent policy learning. One hypothesis is that specifying the goal enables the bottleneck to perform goal-conditioned implicit planning. The planned future may provide the correct waypoints for the agent to reach the distant goal. Figure 6 illustrates a qualitative result of the latent future in AntMaze, which depicts the future states decoded from the bottleneck using the pre-trained Transformer decoder. It demonstrates that the latent future with goal conditioning helps point out the correct direction towards the target location.
Table 3: Ablation study of goal conditioning on AntMaze-Large. Removing the goal conditioning from TrajNet would seriously affect predictive coding objectives and harm the resulting policy performance, suggesting that the properly encoded future representation provides crucial guidance for performing long-horizon tasks. "GC" refers to goal conditioning in TrajNet.
Large-Play	AE-H	MAE-H	MAE-F	MAE-RC	MAE-ALL
w/o GC	32.8 ± 4.3	33.2 ± 11.9	10.0 ± 2.2	16.0 ± 5.3	36.2 ± 10.4
w/ GC	30.2 ± 6.6	36.6 ± 12.6	76.2 ± 4.0	78.2 ± 3.2	32.0 ± 13.2
Goal Location Agent Location
Figure 6: Latent future visualization on AntMaze-Large. A qualitative result comparing the latent future with and without goal conditioning. Given the goal information, the bottleneck can encode the latent future moving in the desired direction.
4.4 Latent Future versus Explicit Future as Policy Conditioning Variables
Prior work [27, 26] has demonstrated that planning into the future is helpful for solving long-horizon tasks. These work performs planning on the explicit future by sampling desirable future states (or transitions), while GCPC leverages goal-conditioned latent representations that encode the future state sequence. In this experiment, we examine how the implicit encoding of future information affects policy performance when it serves as a conditioning variable of PolicyNet. Specifically, we obtain the bottleneck from a TrajNet, which is pre-trained with  $p$ -length future window. The bottleneck encodes latent representations of  $p$  future states. Correspondingly, we acquire the same number of explicit future states using the pre-trained transformer decoder (see Appendix A.4). Either the bottleneck or explicit future states are taken as auxiliary inputs to the PolicyNet. In Table 4, we evaluate the policy with  $p = 70$  in large AntMaze and  $p = 30$  in Kitchen. The results illustrate that the bottleneck is a powerful future information carrier that effectively improves the policy performance. Different from previous approaches that estimate explicit future states with step-by-step dynamics models [27], using latent future representations can mitigate compounding rollout errors. Compared to diffusion-based planning methods [2], which require iterative refinement to obtain a future sequence, latent future avoids the time-consuming denoising step and reduces the decision latency.
Table 4: Comparison between latent and explicit future. Compared to explicit future states, the latent future encoded by the bottleneck is more effective for policy learning.  $p$  is future window size.
Large-Play	Large-Diverse	Kitchen-Mixed	Kitchen-Partial
p = 70	p = 30
Explicit Future	67.9 ± 9.3	70.0 ± 4.6	72.4 ± 4.5	73.9 ± 9.3
Latent Future	78.2 ± 3.2	80.6 ± 3.9	75.6 ± 0.8	90.2 ± 6.6
Table 5: Average normalized scores of GCPC against other baselines on AntMaze. TAP's performances are taken from the original paper [28]. For other baselines, we obtain their performance by re-running author-provided or our replicated implementations with our evaluation protocol. Following [29], we bold all scores within 5 percent of the maximum per task ( $\geq 0.95 \cdot \max$ ).
Dataset	BC	CQL	IQL	DT	TAP	WGCSL	GCIQL	DWSL	RvS-G	GCPC
Umaze	63.4 ± 9.4	88.2 ± 2.3	92.8 ± 3.4	55.6 ± 6.3	-	90.8 ± 2.8	91.6 ± 4.0	71.2 ± 4.2	70.4 ± 4.0	71.2 ± 1.3
Umaze-Diverse	63.4 ± 4.4	47.4 ± 2.0	71.2 ± 7.0	53.4 ± 8.6	-	55.6 ± 15.7	88.8 ± 2.2	74.6 ± 2.8	66.2 ± 5.6	71.2 ± 6.6
Medium-Play	0.6 ± 0.5	72.8 ± 5.7	75.8 ± 1.3	0	78.0	63.2 ± 13.7	82.6 ± 5.4	77.6 ± 3.0	71.8 ± 4.7	70.8 ± 3.4
Medium-Diverse	0.4 ± 0.5	70.8 ± 10.3	76.6 ± 4.2	0	85.0	46.0 ± 12.6	76.2 ± 6.3	74.8 ± 9.3	72.0 ± 3.7	72.2 ± 3.4
Large-Play	0	36.4 ± 10.3	50.0 ± 9.7	0	74.0	0.6 ± 1.3	40.0 ± 16.2	15.2 ± 7.7	35.6 ± 7.6	78.2 ± 3.2
Large-Diverse	0	36.0 ± 8.3	52.6 ± 5.9	0	82.0	2.4 ± 4.3	29.8 ± 6.8	19.0 ± 2.8	25.2 ± 4.8	80.6 ± 3.9
Ultra-Play	0	18.0 ± 13.3	21.2 ± 7.5	0	22.0	0.2 ± 0.4	20.6 ± 7.6	25.2 ± 3.0	25.6 ± 6.7	56.6 ± 9.5
Ultra-Diverse	0	9.6 ± 14.6	17.8 ± 4.0	0	26.0	0	28.4 ± 11.8	25.0 ± 8.6	26.4 ± 7.7	54.6 ± 10.3
Average	16.0	47.4	57.3	13.6	-	32.4	57.3	47.8	49.2	69.4
4.5 Effectiveness of GCPC
Finally, we show the effectiveness of GCPC by evaluating it on three different domains: AntMaze, Kitchen and Gym Locomotion.
Baselines and prior methods. We compare our approach to both supervised learning methods and value-based RL methods. For the former, we consider: (1) Behavioral Cloning (BC), (2) RvS-R/G [15], a conditional imitation learning method that is conditioned on either a target state or an expected return. (3) Decision Transformer (DT) [10], a return-conditioned model-free method that learns a transformer-based policy, (3) Trajectory Transformer (TT) [27], a model-based method that performs beam search on a transformer-based trajectory model for planning, and (4) Decision Diffuser (DD) [2], a diffusion-based planning method that synthesizes future states with a diffusion model and acts by computing the inverse dynamics. For the latter, we select methods based on dynamic programming, including (5) CQL [31] and (6) IQL [29]. Additionally, we include three goal (state)-conditioned baselines for AntMaze and Kitchen: (7) Goal-conditioned IQL (GCIQL), (8) WGCSL [52] and (9) DWSL [25]. AntMaze-Ultra is a customized environment proposed by [28], therefore we include TAP's performance on the v0 version of AntMaze for completeness. When feasible, we re-run the baselines with our evaluation protocol for fair comparison [1].
Our empirical results on AntMaze and Kitchen are presented in Table 5 and Table 6. We find our methods outperform all previous methods on the average performance. In particular, on the most challenging large and ultra AntMaze environments, our method achieves significant improvements over the RvS-G baseline, demonstrating the efficacy of learning good future representations using sequence modeling in long-horizon tasks. Table 7 shows the results on Gym Locomotion tasks. Our approach obtains competitive performance as prior methods. We also notice that compared to Decision Transformer, RvS-R can already achieve strong average performance. This suggests that for some tasks sequence modeling may not be a necessary component for policy improvement. With a large fraction of near-optimal trajectories in the dataset, a simple MLP policy may provide enough capacity to handle most of the locomotion tasks.
Table 6: Average normalized scores of GCPC against other baselines on Kitchen. CQL and DD's results are taken from [2]. For other baselines, we obtain their performance by re-running authorprovided or our replicated implementations with our evaluation protocol. Following [29], we bold all scores within 5 percent of the maximum per task ( $\geq 0.95 \cdot \max$ ).
Dataset	BC	CQL	IQL	DT	DD	WGCSL	GCIQL	DWSL	RvS-G	GCPC
Mixed	48.9 ± 0.7	52.4	53.2 ± 1.6	50.7 ± 7.1	65.0	77.8 ± 3.6	74.6 ± 1.9	74.6 ± 0.6	69.4 ± 4.2	75.6 ± 0.8
Partial	41.3 ± 3.7	50.1	59.7 ± 8.3	48.6 ± 9.5	57.0	75.2 ± 6.4	74.7 ± 4.1	74.0 ± 5.8	71.7 ± 7.9	90.2 ± 6.6
Average	45.1	51.3	56.5	49.7	61.0	76.5	74.7	74.3	70.6	82.9
Table 7: Average normalized scores of our approach against other baselines on Gym Locomotion. TT and DD's results are taken from [2]. For other baselines, we obtain their performance by re-running author-provided or our replicated implementations with our evaluation protocol. Following [29], we bold all scores within 5 percent of the maximum per task ( $\geq 0.95 \cdot \max$ ).
Dataset	Environment	BC	CQL	IQL	DT	TT	DD	RvS-R	GCPC
Medium-Expert	HalfCheetah	61.9 ± 5.8	87.7 ± 7.0	92.4 ± 0.4	88.8 ± 2.6	95	90.6	93.4 ± 0.3	94.0 ± 0.9
Medium-Expert	Hopper	55.1 ± 2.8	110.5 ± 2.9	104.2 ± 8.7	108.4 ± 2.0	110	111.8	111.3 ± 0.2	111.7 ± 0.4
Medium-Expert	Walker2d	100.4 ± 13.4	110.4 ± 0.6	110.2 ± 0.7	108.6 ± 0.3	101.9	108.8	109.6 ± 0.4	109.0 ± 0.2
Medium	HalfCheetah	43.0 ± 0.4	47.1 ± 0.3	47.7 ± 0.2	42.9 ± 0.2	46.9	49.1	44.2 ± 0.2	44.5 ± 0.5
Medium	Hopper	55.8 ± 2.8	70.1 ± 1.6	69.2 ± 3.2	67.8 ± 4.2	61.1	79.3	65.1 ± 5.2	68.0 ± 4.4
Medium	Walker2d	74.1 ± 3.2	83.5 ± 0.5	84.5 ± 1.5	76.5 ± 1.6	79	82.5	78.4 ± 2.6	78.0 ± 2.4
Medium-Replay	HalfCheetah	37.2 ± 1.3	45.4 ± 0.3	44.9 ± 0.3	37.8 ± 0.9	41.9	39.3	40.2 ± 0.2	40.7 ± 1.5
Medium-Replay	Hopper	33.7 ± 8.5	96.2 ± 1.9	93.9 ± 9.1	78.0 ± 11.6	91.5	100	88.5 ± 12.9	94.2 ± 3.5
Medium-Replay	Walker2d	19.2 ± 7.3	79.8 ± 1.6	78.6 ± 5.7	72.5 ± 3.3	82.6	75	71.0 ± 5.1	77.6 ± 9.8
Average	53.4	81.2	80.6	75.7	78.9	81.8	78.0	79.7
In this work, we aim to investigate the role of sequence modeling in learning trajectory representations and its utility in acquiring performant policies. To accomplish this, we employ a two-stage framework that decouples trajectory representation learning and policy learning. This framework unifies many existing RvS methods, enabling us to study the impacts of different trajectory representation learning objectives for sequential decision making. Within this framework, we introduce a specific design - Goal-conditioned Predictive Coding (GCPC), that incorporates a compact bottleneck to transfer representations between the two stages and learns goal-conditioned latent representations modeling the future trajectory. Through extensive experiments, we observe that the bottleneck generated by GCPC properly encodes the goal-conditioned future information and brings significant improvement over some long-horizon tasks. By empirically evaluating our approach on three benchmarks, we demonstrate GCPC achieves competitive performance on a variety of tasks.
Limitations and Future work. GCPC models the future by performing maximum likelihood estimation on offline collected trajectories, which may predict overly optimistic future behaviors and lead to suboptimal actions in stochastic environments. Future work includes discovering policies that are robust to the environment stochasticity by considering multiple possible futures generated by GCPC. Another limitation of our work is that GCPC may not be sufficient to maintain high accuracy for long-term future prediction when high-dimensional states are involved, which may potentially be tackled by leveraging foundation models to acquire representations for high-dimensional inputs.
We appreciate all anonymous reviewers for their constructive feedback. We would like to thank Calvin Luo and Haotian Fu for their discussions and insights, and Tian Yun for the help on this project. This work is in part supported by Adobe, Honda Research Institute, Meta AI, Samsung Advanced Institute of Technology, and a Richard B. Salomon Faculty Research Award for C.S.
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In Advances in neural information processing systems, 2021.
[2] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In International Conference on Learning Representations, 2023.
[3] Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. In Advances in Neural Information Processing Systems, 2022.
[4] Hangbo Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. In International Conference on Learning Representations, 2022.
[5] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 2013.
[6] Prajjwal Bhargava, Rohan Chitnis, Alborz Geramifard, Shagun Sodhani, and Amy Zhang. Sequence modeling is a robust contender for offline reinforcement learning. arXiv preprint arXiv:2305.14550, 2023.
[7] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? In Advances in Neural Information Processing Systems, 2022.
[8] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in neural information processing systems, 2020.
[9] Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anca Dragan, and Sam Devlin. Uni[MASK]: Unified inference in sequential decision problems. In Advances in Neural Information Processing Systems, 2022.
[10] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345, 2021.
[11] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[12] Yiming Ding, Ignasi Clavera, and Pieter Abbeel. Mutual information maximization for robust plannable representations. arXiv preprint arXiv:2005.08114, 2020.
[13] Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Philipp. Goal-conditioned imitation learning. In Advances in neural information processing systems, volume 32, 2019.
[14] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
[49] Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. Masked trajectory models for prediction, representation, and control. In International Conference on Machine Learning, 2023.
[50] Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173, 2022.
[51] Mengjiao Yang and Ofir Nachum. Representation matters: offline pretraining for sequential decision making. In International Conference on Machine Learning. PMLR, 2021.
[52] Rui Yang, Yiming Lu, Wenzhe Li, Hao Sun, Meng Fang, Yali Du, Xiu Li, Lei Han, and Chongjie Zhang. Rethinking goal-conditioned supervised learning and its connection to offline RL. In International Conference on Learning Representations, 2022.
A.1 Additional Baseline Details
We provide additional information on the baseline methods and their performance we present in the main paper. We compare our proposed method with (1) supervised learning methods, including Behavioral Cloning (BC), RvS, Decision Transformer (DT), Trajectory Transformer (TT), and Decision Diffuser (DD). We also compare with (2) value-based methods, including Conservative Q-Learning (CQL) and Implicit Q-Learning (IQL). DT uses a transformer encoder to model offline-RL as an autoregressive sequence generation problem in a model-free way, which predicts future actions based on previous states, actions, and returns-to-go. Similarly, The model-based TT also adopts transformer structure for sequence prediction but uses beam search for planning during execution. Inspired by conditional generative model, DD adopts a diffusion model to sample future state sequence based on returns or constraints and extract actions between the states by using a inverse dynamics model with relying on value estimation. Similar to our PolicyNet, RvS-G and RvS-R take the current state and the goal (return or state) as input, and train a two-layer MLP to predict the actions via supervised learning. For value-based methods, in order to solve the problem of overestimating values induced by distributional shift, CQL adds value regularization terms to the standard Bellman error objective to learn a lower bound of the true Q-function. IQL approximates the policy improvement step implicitly instead of updating the Q-function with target actions sampled from the behavior policy. For the additional goal (state)-conditioned baselines, GCIQL is a goal-conditioned version of IQL. WGCSL enhances GCSL with an advanced compound weight, which optimizes the lower-bound of the goal-conditioned RL objective. DWSL first models the distance between states and the goal, and extracts the policy by imitating the actions that reduce the minimum distance metric.
Baseline Performance Source. For AntMaze, we use v2 version for D4RL environments and v0 version for the additional AntMaze-Ultra environments. For Gym Locomotion and Kitchen, we use v2 and v0 version of the datasets from the D4RL benchmark, respectively. For all baselines except TT, DD and TAP, we obtain their performance by re-running author-provided (CQL, IQL, DT, WGCSL, GCIQL, DWSL) or our replicated (BC, RvS-G/R) implementations with our evaluation protocol. For CQL's performance on Kitchen, we are unable to reproduce the previously reported performance with the default hyperparameters, so we take the corresponding results from [2]. For additional goal (state)-conditioned baselines (i.e. WGCSL, GCIQL and DWSL), we take the implementations provided by DWSL authors<sup>1</sup>, and use the same goal specification (see details in A.6) for AntMaze and Kitchen as ours. For TT, DD and TAP's performance, we take the results from the original papers.
A.2 Evaluation Protocol