SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zilai Zeng
Ce Zhang
Shijie Wang
Recent work has demonstrated the effectiveness of formulating decision making as supervised learning on offline-collected trajectories. Powerful sequence models, such as GPT or BERT, are often employed to encode the trajectories. However, the benefits of performing sequence modeling on trajectory data remain unclear. In this work, we investigate whether sequence modeling has the ability to condense trajectories into useful representations that enhance policy learning. We adopt a two-stage framework that first leverages sequence models to encode trajectory-level representations, and then learns a goal-conditioned policy employing the encoded representations as its input. This formulation allows us to consider many existing supervised offline RL methods as specific instances of our framework. Within this framework, we introduce Goal-Conditioned Predictive Coding (GCPC), a sequence modeling objective that yields powerful trajectory representations and leads to performant policies. Through extensive empirical evaluations on AntMaze, FrankaKitchen and Locomotion environments, we observe that sequence modeling can have a significant impact on challenging decision making tasks. Furthermore, we demonstrate that GCPC learns a goal-conditioned latent representation encoding the future trajectory, which enables competitive performance on all three benchmarks. Our code is available at https://brown-palm.github.io/GCPC/.
Goal-conditioned imitation learning [13, 15, 20] has recently emerged as a promising approach to solve offline reinforcement learning problems. Instead of relying on value-based methods, they directly learn a policy that maps states and goals (e.g. expected returns, or target states) to the actions. This is achieved by supervised learning on offline collected trajectories (i.e. sequences of state, action, and reward triplets). This learning paradigm enables the resulting framework to avoid bootstrapping for reward propagation, a member of the "Deadly Triad" [45] of RL which is known to lead to unstable optimization when combining with function approximation and off-policy learning, and also allows the model to leverage large amounts of collected trajectories and demonstrations. The emergence of RL as supervised learning on trajectory data coincides with the recent success of the Transformer architecture [46, 20] and its applications to sequence modeling, such as GPT [8] for natural language and VPT [3] for videos. Indeed, several recent works have demonstrated the effectiveness of sequence modeling for offline RL [10, 27, 34, 44, 49, 9, 42]. Most of these approaches apply the Transformer architecture to jointly learn the trajectory representation and the policy. When the sequence models are unrolled into the future, they can also serve as "world models" [21, 22] for many RL applications. By leveraging advanced sequence modeling objectives from language and visual domains [11, 4, 24], such methods have shown competitive performance in various challenging tasks [16, 5].
Despite the enthusiasm and progress, the necessity of sequence modeling for offline reinforcement learning has been questioned by recent work [15]: With a simple but properly tuned neural architecture (e.g. a multilayer perceptron), the trained agents can achieve competitive performance on several
challenging tasks while taking only the current state and the overall goal as model inputs. In some tasks, these agents even significantly outperform their counterparts based on sequence modeling. These observations naturally motivate the question: Is explicit sequence modeling of trajectory data necessary for offline RL? And if so, how should it be performed and utilized?
To properly study the impact of sequence modeling for decision making, we propose to decouple trajectory representation learning and policy learning. We adopt a two-stage framework, where sequence modeling can be applied to learn the trajectory representation, the policy, or both. The connection between the two stages can be established by leveraging encoded trajectory representations for policy learning, or by transferring the model architecture and weights (i.e. the policy can be initialized with the same pre-trained model used for trajectory representation learning). These stages can be trained separately, or jointly and end-to-end. This design not only facilitates analyzing the impacts of sequence modeling, but is also general such that prior methods can be considered as specific instances, including those that perform joint policy learning and trajectory representation learning [10, 34, 49], and those that learn a policy directly without sequence modeling [15].
Concretely, we aim at investigating the following questions: (1) Are offline trajectories helpful due to sequence modeling, or simply by providing more data for supervised policy learning? (2) What would be the most effective trajectory representation learning objectives to support policy learning? Should the sequence models learn to encode history experiences [10], future dynamics [27], or both [19]? (3) As the same sequence modeling framework may be employed for both trajectory representation learning and policy learning [34, 9, 49], should they share the same training objectives or not?
We design and conduct experiments on the AntMaze, FrankaKitchen, and Locomotion environments to answer these questions. We observe that sequence modeling, if properly designed, can effectively aid decision-making when its resulting trajectory representation is used as an input for policy learning. We also find that there is a discrepancy between the optimal self-supervised objective for trajectory representation learning, and that for policy learning. These observations motivate us to propose a specific design of the two-stage framework: It compresses trajectory information into compact bottlenecks via sequence modeling pre-training. The condensed representation is then used for policy learning with a simple MLP-based policy network. We observe that goal-conditioned predictive coding (GCPC) is the most effective trajectory representation learning objective. It enables competitive performance across all benchmarks, particularly for long-horizon tasks. We attribute the strong empirical performance of GCPC to the acquisition of goal-conditioned latent representations about the future, which provide crucial guidance for decision making.
To summarize, our main contributions include:
- We propose to decouple sequence modeling for decision making into a two-stage framework, namely trajectory representation learning and policy learning. It provides a unified view for several recent reinforcement learning via supervised learning methods.
- We conduct a principled empirical exploration with our framework to understand if and when sequence modeling of offline trajectory data benefits policy learning.
- We discover that goal-conditioned predictive coding (GCPC) serves as the most effective sequence modeling objective to support policy learning. Our overall framework achieves competitive performance on AntMaze, FrankaKitchen and Gym Locomotion benchmarks.
What is essential for Offline RL? Offline reinforcement learning aims to obtain effective policies by leveraging previously collected datasets. Prior work usually adopts dynamic programming [31, 29, 18] or supervised behavioral cloning (BC) methods [43, 3]. Recent approaches [10, 27] demonstrate the effectiveness of solving decision making tasks with sequence modeling, whereas RvS [15] establishes a strong MLP baseline for conditional behavioral cloning. To further understand what are the essential components for policy learning, researchers have investigated the assumptions required to guarantee the optimality of return-conditioned supervised learning [7], and examined offline RL agents from three fundamental aspects: representations, value functions, and policies [17]. Another concurrent work [6] compares the preferred conditions (e.g. data, task, and environments) to perform Q-learning and imitation learning. In our work, we seek to understand how sequence modeling may benefit offline RL and study the impacts of different sequence modeling objectives.
Transformer for sequential decision-making. Sequential decision making has been the subject of extensive research over the years. The tremendous success of the Transformer model for natural language processing [11, 40] and computer vision [24, 35] has inspired numerous works that seek to apply such architectures for decision making, and similarly motivates our work. Prior work has shown how to model sequential decision making as autoregressive sequence generation problem to produce a desired trajectory [10], while others have explored the applications of Transformer for model-based RL [27, 47] and multi-task learning [41, 9]. Our work aims to utilize Transformer-based models to learn good trajectory representations that can benefit policy learning.
Masked Autoencoding. Recent work in NLP and CV has demonstrated masked autoencoding (MAE) as an effective task for self-supervised representation learning [11, 8, 4, 24]. Inspired by this, Uni[MASK] [9], MaskDP [34], and MTM [49] propose to train unified masked autoencoders on trajectory data by randomly masking the states and actions to be reconstructed. These models can be directly used to solve a range of decision making tasks (e.g. return-conditioned BC, forward dynamics, inverse dynamics, etc.) by varying the masking patterns at inference time, without relying on task-specific fine-tuning. In contrast, we decouple trajectory representation learning and policy learning into two stages. Unlike approaches which perform value-based offline RL (e.g. TD3, Actor-Critic) after pretraining [34, 49], our policy learning stage adopts imitation learning setup where the policy can be learned without maximizing cumulative reward.
Self-supervised Learning for RL. Self-supervised learning has emerged as a powerful approach for learning useful representations in various domains, including reinforcement learning (RL). In RL, self-supervised learning techniques aim to leverage unlabeled or partially observed data to pre-train agents or learn representations that facilitate downstream RL tasks. Prior work mainly focuses on using self-supervised learning to get state representations [51, 38, 36, 50, 10, 34, 9, 49, 32] or world models [22, 23, 44, 27, 12, 37]. In [51], the authors evaluate a broad set of state representation learning objectives on offline datasets and demonstrate the effectiveness of contrastive self-prediction. In this work, we investigate representation learning objectives in trajectory-space with the sequence modeling tool, which enables us to explore the impact of different modalities (e.g. states, actions, goals, etc.) on trajectory representation learning.
We revisit the role of sequence modeling in offline reinforcement learning, from the perspectives of trajectory representation learning and policy learning. Section 3.1 describes the background knowledge on extracting policies from offline trajectories using supervised learning. Section 3.2 introduces a two-stage framework that decouples trajectory representation learning and policy learning, which serves as the basis for our investigation. In Section 3.3, we propose a specific instantiation of the two-stage framework, in which we leverage a self-supervised learning objective - Goal-Conditioned Predictive Coding (GCPC), to acquire performant trajectory representations for policy learning.
3.1 Offline Reinforcement Learning via Supervised Learning
We follow prior work that leverages sequence modeling for offline reinforcement learning, and adopt the Reinforcement learning via Supervised learning (RvS) (e.g. [15]) setting. RvS aims to solve the offline RL problem as conditional, filtered, or weighted imitation learning. It assumes that a dataset has been collected offline, but the policies used to collect the dataset may be unknown, such as with human experts or a random policy. The dataset contains a set of trajectories:  $\mathcal{D} = \{\tau^i\}_{i=1}^N$ . Each trajectory  $\tau^i$  in the dataset is represented as  $\{(s_t, a_t)\}_{t=1}^H$ , where  $H$  is the length of the trajectory, and  $s_t, a_t$  refer to the state, action at timestep  $t$ , respectively. A trajectory may optionally contain the reward  $r_t$  received at timestep  $t$ .
As the trajectories are collected with unknown policies, they may not be optimal or have expert-level performance. Prior work [10, 33, 30] has shown that properly utilizing offline trajectories containing suboptimal data might produce better policies. Intuitively, the suboptimal trajectories may still contain sub-trajectories that demonstrate useful "skills", which can be composed to solve new tasks.
Given the above trajectories, we denote a goal-conditioned policy  $\pi_{\theta}$  as a function parameterized by  $\theta$  that maps an observed trajectory  $\tau_{\mathrm{obs}}^{(t)} = \{s_{\leq t},a_{< t}\}$  and a goal  $g^{(t)}$  to an action  $a_{t}$ . The goal  $g^{(t)}$  is computed from  $\tau_{t:H}$ , which can be represented as either a target state configuration sampled
from  $s_{t:H}$  or an expected cumulative reward/return-to-go computed from  $r_{t:H}$  (see Appendix A.6 for details). For simplicity of notation, we write  $\tau_{\mathrm{obs}}^{(t)}$  as  $\tau_{\mathrm{obs}}$  and  $g^{(t)}$  as  $g$ . We consider a policy should be able to take any form of state or trajectory information as input and predict the next action: (1) when only the current observed state  $s_t$  and the goal  $g$  are used, the policy  $\hat{a}_t = \pi_\theta(s_t, g)$  ignores the history observations; (2) when  $\pi_\theta$  is a sequence model, it can employ the whole observed trajectory to predict the next action  $\hat{a}_t = \pi_\theta(\tau_{\mathrm{obs}}, g)$ . To optimize the policy, a commonly used objective is to find the parameters that fit the mapping of observations to actions using maximum likelihood estimation:
$$
\theta^ {*} = \underset {\theta} {\operatorname {a r g m a x}} \mathbb {E} _ {\tau \sim \mathcal {D}} \left[ \prod_ {t = 1} ^ {| \tau |} \pi_ {\theta} \left(a _ {t} \mid \tau_ {\mathrm {o b s}}, g\right) \right] \tag {1}
$$
3.2 Decoupled Trajectory Representation and Policy Learning
Sequence modeling can be used for decision making from two perspectives, namely trajectory representation learning and policy learning. The former aims to acquire useful representations from raw trajectory inputs, often in the form of a condensed latent representation, or the pretrained network weights themselves. The latter aims to map the observations and the goal into actions that accomplish the task. To explicitly express the trajectory representation function, we rewrite the goal-conditioned policy function as follows:
$$
\hat {a} _ {t} = \pi_ {\theta} (f (\tau_ {\mathrm {o b s}}), g) \tag {2}
$$
where  $f(\cdot)$  is an arbitrary function, such as a neural network, that computes representation from  $\tau_{\mathrm{obs}}$ .  $f(\cdot)$  can also be an identity mapping such that  $\pi_{\theta}(\cdot)$  directly operates on  $\tau_{\mathrm{obs}}$ .
Motivated by the recent success of sequence modeling in NLP [11, 40] and CV [24, 14], both the trajectory learning function and the policy learning function can be implemented with Transformer neural networks as  $f_{\phi}(\cdot)$  and  $\pi_{\theta}(\cdot)$ . We hypothesize that it is beneficial for  $f_{\phi}$  to condense the trajectories into a compact representation using sequence modeling techniques. We also hypothesize that it is desirable to decouple the trajectory representation learning from policy learning. The decoupling not only offers flexibility on the choice of representation learning objectives, but also allows us to study the impact of sequence modeling for trajectory representation learning and policy learning independently. We thus adopt a two-stage framework with a TrajNet  $f_{\phi}(\cdot)$  and a PolicyNet  $\pi_{\theta}(\cdot)$ . TrajNet aims to learn trajectory representations with self-supervised sequence modeling objectives, such as masked autoencoding [24] or next token prediction [40]. PolicyNet aims to obtain a performant policy with the supervised learning objective in Equation 1.
We now describe an overall flow of the two networks, which we will show is general to represent recently proposed methods. In the first stage, we approach trajectory representation learning as masked autoencoding. TrajNet receives a trajectory  $\tau$  and an optional goal  $g$ , and is trained to reconstruct  $\tau$  from a masked view of the same. Optionally, TrajNet also generates a condensed trajectory representation  $B$ , which can be utilized by PolicyNet for subsequent policy learning:
$$
\hat {\tau}, B = f _ {\phi} (\operatorname {M a s k e d} (\tau), g) \tag {3}
$$
During the second stage, TrajNet is applied on the unmasked observed trajectory  $f_{\phi}(\tau_{\mathrm{obs}}, g)$ , to obtain  $B_{\mathrm{obs}}$ . PolicyNet then predicts the action  $a$  given  $\tau_{\mathrm{obs}}$  (or the current observed state  $s_t$ ), the goal  $g$ , and trajectory representation  $B_{\mathrm{obs}}$ :
$$
a = \pi_ {\theta} \left(\tau_ {\mathrm {o b s}}, B _ {\mathrm {o b s}}, g\right) \tag {4}
$$
Our framework provides a unified view to compare different design choices (e.g. input information, architecture, training objectives, etc.) for representation learning and policy learning, respectively. Many existing methods can be seen as special cases of our framework. For example, prior works [34, 49] instantiate TrajNet with a bi-directional Transformer pre-trained via masked prediction, and re-use it as the PolicyNet in a zero-shot manner. To implement DT [10],  $f(\cdot)$  is set as an identity mapping function of the input trajectory and  $\pi (\cdot)$  is trained to autoregressively generate actions. These methods learn trajectory representations and the policy jointly, with training objectives inspired by MAE [24] and GPT [40]. At last, our framework recovers RvS-G/R [15] by having the output of  $f(\cdot)$  in Eq. 2 as the last observed state  $s_t$  from  $\tau_{\mathrm{obs}}$ .
Figure 1: Stage 1 - Trajectory Representation Learning. For notation consistency between the two stages, we separate the input into observed (history) states and future states. Left: TrajNet. We input randomly masked history state, goal and slot tokens to the transformer encoder. The decoder takes in encoded slot tokens (the bottleneck) and a sequence of masked tokens, and reconstruct the whole trajectory. With this training objective, we encourage the bottleneck to perform predictive coding, which is conditioned on the goal and history states. Right: All Masking Patterns. TrajNet can be trained with different masking patterns with their corresponding reconstruction objectives. The illustration of TrajNet on the left uses "MAE-RC", which is adopted by GCPC.
3.3 Goal-Conditioned Predictive Coding
Finally, we introduce a specific design of the two-stage framework, Goal-Conditioned Predictive Coding (GCPC), which is inspired by our hypotheses introduced in Section 3.2. To facilitate the transfer of trajectory representations between the two stages, we compress the trajectory into latent representations using sequence modeling, which we refer to as the bottleneck. With this design, we train the bottleneck to perform goal-conditioned future prediction, so that the latent representations encode future behaviors toward the desired goal, which are subsequently used to guide policy learning.
In the first stage (as shown in Figure 1), we use a bi-directional Transformer as TrajNet. The inputs include  $T$  state tokens, one goal token, and a few learnable slot tokens. The action tokens are ignored. We first embed the goal token and all state tokens with separate linear encoders, then apply sinusoidal positional encoding before all inputs are sent into the Transformer Encoder.
Figure 2: Stage 2 - Policy Learning. We implement policy learning with a simple MLP as PolicyNet. We input unmasked history states and retrieve the bottleneck from pre-trained encoder. Then the bottleneck is taken as the input to the policy network. We view the bottleneck generated by pre-trained encoder as goal-conditioned latent representations for the future trajectory.
While the full trajectories can be used for trajectory representation learning with TrajNet, only observed states are available during policy learning. For notation consistency, we denote the  $k$  observed states as "history", and the  $p$  states that follow the observed states as "future". The total number of input states for stage 1 is  $T = k + p$ , and for stage 2 is  $k$ . Both  $k$  and  $p$  are hyperparameters that represent history window length and future window length, respectively.
To perform goal-conditioned predictive coding, the entire future trajectory is masked. The tokens in the observed history are randomly masked (Figure 1 "MAE-RC"). The bottleneck is taken as the encoded slot tokens, which condenses the masked input trajectory into a compact representation. The decoder takes the bottleneck and  $T$  masked tokens as the input, and aims to reconstruct both the history and the future states. Our intuition is that this will encourage the bottleneck to learn representations encoding the future trajectory and provide useful signals for decision making. The training objective of the TrajNet is to minimize MSE loss between the reconstructed and the ground-truth trajectory.
Figure 2 illustrates the interface between the TrajNet and the PolicyNet. PolicyNet is implemented as a simple MLP network. The trained TrajNet has its weights frozen, and only the TrajNet encoder is used to compute the trajectory representation  $B_{\mathrm{obs}}$  from  $k$  unmasked history states. PolicyNet takes the current state  $s_t$ , goal  $g$ , and bottleneck  $B_{\mathrm{obs}}$  as input, and outputs an action. The training objective of PolicyNet is to minimize MSE loss between the predicted and ground-truth actions.
Discussion. Decoupling trajectory representation learning and policy learning allows us to explore different model's inputs and sequence modeling objectives. For example, recent work [2, 39] observed that the action sequences could potentially be detrimental to the learned policy in some tasks. In GCPC, we employ state-only trajectories as input and ignore the actions. Different objectives (e.g. masking patterns in the masked autoencoding objective) also have an impact on what is encoded in the bottleneck. Here we introduce five sequence modeling objectives for trajectory representation learning (as shown in Table 1) and study their impacts on policy learning. When using "AE-H" or "MAE-H", the bottleneck is only motivated to summarize the history states. When using "MAE-F" or "MAE-RC", the bottleneck is asked to perform predictive coding, and thus encode the future state sequences to achieve the provided goal. By default, we adopt the "MAE-RC" objective in GCPC.
Table 1: Objectives for Trajectory Representation Learning
Input	Reconstruct
History	Future	History	Future
AE-H	Unmasked	-	✓
MAE-H	Randomly Masked	-	✓
MAE-F	Unmasked	Fully Masked	✓	✓
MAE-RC	Randomly Masked	Fully Masked	✓	✓
MAE-ALL	Randomly Masked	Randomly Masked	✓	✓
In this section, we aim to answer the following questions with empirical experiments: (1) Does sequence modeling benefit reinforcement learning via supervised learning on trajectory data, and how? (2) Is it beneficial to decouple trajectory representation learning and policy learning with bottlenecks? (3) What are the most effective trajectory representation learning objectives?