CLaP: Conditional Latent Planners for Offline Reinforcement Learning
Abstract
Recent work has formulated offline reinforcement learning (RL) as a sequence modeling problem, benefiting from the simplicity and scalability of the Transformer architecture.
However, sequence models struggle to model trajectories that are long-horizon or involve complicated environment dynamics.
We propose CLaP (Conditional Latent Planners) to learn a simple goal-conditioned latent space from offline agent behavior, and incrementally decode good actions from a latent plan.
We evaluate our method on continuous control domains from the D4RL benchmark.
Compared to non-sequential models and return-conditioned sequential models, CLaP shows competitive if not better performance across continuous control tasks.
It particularly does better in environments with complex transition dynamics with up to +149 .
8% performance increase.
Our results suggest that decision-making is easier with simplified latent dynamics that models behavior as being goal-conditioned.
Introduction
Reinforcement learning (RL) has seen significant advances with the introduction of deep neural networks, achieving excellent results across a wide range of domains (Tesauro, 1994; Hafner and
Riedmiller, 2011; Levine and Koltun, 2013; Mnih et al., 2013; Levine et al., 2016; Silver et al., 2017; Kalashnikov et al., 2018).
These RL methods generally involve continued interactions with the environment, using experience to improve the learned policy.
However, online interaction is expensive and potentially dangerous.
This prompts the need for offline RL methods that allow for data to be collected in a safe way and then have agents learn from this offline dataset (Haarnoja et al., 2018; Fujimoto et al., 2019; Wu et al., 2019; Kumar et al., 2020).
Learning from offline data is challenging because the pre-collected dataset of agent behavior may be limited in quantity and quality (Levine et al., 2020).
Low data quantity leads to generalization issues, where a majority of states are out-of-distribution from the agent’s training distribution.
Low data quality leads to learning sub-optimal behavior, where the provided demonstrations are not generated by experts.
Recent works apply sequence models for decision making to overcome these issues (Janner et al., 2021;
Chen et al., 2021).
Sequence models streamline learning by not requiring additional algorithm constraints such as policy regularization (Fujimoto et al., 2019; Kumar et al., 2019) and conservatism (Yu et al., 2020; Kidambi et al., 2020).
However, there are two limitations in prior sequence modeling approaches for RL.
Chen et al.
(2021) models trajectories as return-conditioned, which can cause issues when the optimal return is not known (Brandfonbrener et al., 2022).
Janner et al.
(2021)’s approach models trajectories as goal-conditioned by prepending the goal at the start of the sequence.
However, sequence models struggle to attend to goals when the sequences are long or the environment dynamics is complicated (Kiddon et al., 2016; Fan et al., 2019; Hua and Wang, 2020).
Recent work in natural language processing (NLP) overcome these issues by modeling text as generated from a goal-conditioned latent process (Wang et al., 2021).
The work assumes the latent dynamics is simple and known , which makes generating long coherent sequences easier.
Our work proposes Conditional Latent Planners (CLaP), a sequence-based model that decodes actions from a learned goal-conditioned latent space.
We incorporate the ideas from Wang et al.
(2021), which generates language from a goal-conditioned latent process, and apply them to the offline
RL setting.
In the offline RL setting, CLaP generates actions as sampled from a goal-conditioned latent process learned from offline data.
Note that despite methodological similarities, the offline
RL setting is very different from the NLP setting.
For example, they assume the sequence model has been pretrained and can be further finetuned on large datasets, whereas we assume access to smaller and potentially low-quality data.
We evaluate our method against a prior state-of-the art non-sequential method Conservative Q-Learning and the original Decision Transformer model on the D4RL benchmark (Fu et al., 2020; Kumar et al., 2020; Chen et al., 2021).
We show that CLaP outperforms Conservative Q-Learning and Decision Transformer on tasks with particularly complex transition dynamics and matches their performance on other tasks.
Related Work
Latent planning methods Planning in unknown environments is difficult because it requires a faithful model of the world.
Obtaining this model is often challenging and even impossible, especially in high-dimensional or image-based domains (Hafner et al., 2019; Moerland et al., 2020; Schrittwieser et al., 2021).
Prior work circumvent the need to specify a world model for the raw state space with latent-based methods (Ha and Schmidhuber, 2018; Oord et al., 2018).
They learn the latent space with a simple Variational Autoencoder or noise-based contrastive learning (Gutmann and Hyvärinen, 2010; Kingma and Welling, 2013; Rezende et al., 2014).
However, without assuming temporal structure in the latent space, these methods often fail to predict future latents accurately.
Accumulated prediction errors can result in decoding bad actions, leading to low returns.
Our method overcomes these challenges by assuming a goal-conditioned structure for planning in latent space.
Wang et al.
(2021) is a recent work in NLP which models language with goal-conditioned stochastic processes, namely the Brownian bridge process (Revuz and Yor, 2013).
Their work assumes that each document is generated from a process that is pinned to the same latent start and goal.
Intermediate latents are Gaussian distributions centered at interpolated points between the two, and with variances that are a function of time.
Despite methodological similarities, we note key differences in the setting.
For example, their work assumes that the sequence model has been pretrained on a large dataset (Radford et al., 2019) whereas ours is pretrained on an offline dataset that is much smaller in comparison.
Their work assumes that the model has been pretrained and can be further finetuned on high-quality text, whereas we assume a varying quality of data.
Finally, they assume that the documents are generated from a process pinned at the same start and goal, whereas we assume that the trajectories are generated from processes with different start and goal states.
Our work focuses on addressing the challenges that arise from these differences.
Goal-conditioned RL A number of prior works have trained goal-conditioned RL policies using supervised learning and model-free RL (Kaelbling, 1993; Nair et al., 2018; Ghosh et al., 2019;
Eysenbach et al., 2019).
In the offline setting, many methods employ goal-conditioned Q-learning with hindsight relabeling techniques to incorporate goal information and to be sample efficient (Andrychowicz et al., 2017; Chebotar et al., 2021; Tian et al., 2021).
These relabeling methods are related to how CLaP learns a goal-conditioned latent structure over agent trajectories.
However, unlike relabeling methods, CLaP uses this latent structure to simplify dynamics and actively decode intermediate latents and actions needed to reach a desired goal state.
Sequence models for RL Our work builds off of recent sequence-modeling approaches for RL, such as Decision Transformer (Chen et al., 2021) and Trajectory Transformer (Janner et al., 2021).
Decision Transformer is a causally masked Transformer that generates actions conditioned on past states, actions, and desired returns.
Trajectory Transformer is a predictive dynamics model trained on discretized states, actions, and rewards, which behaves like a model-based RL method when paired with planning algorithms.
Trajectory Transformer’s goal-reaching RL scheme prepends goal states to agent states in a raw, discretized format.
However, such a basic form of goal-conditioning does not yield useful information on what goal-conditioning behavior actually looks like, especially for long-horizon tasks in complex environments.
CLaP instead explicitly learns a latent space with goal-conditioned behavior and then uses it to decode good actions in a step-by-step, more granular fashion.
Offline RL Formalism
We formalize the offline RL setting described by the tuple ( S , A , P, R ) , consisting of states s ∈ S , actions a ∈ A , transition dynamics P ( s ′ | s, a ) , and a reward function r =
R ( s, a ) .
Similar to Chen et al.
(2021), we use s t , a t , and r t =
R ( s t , a t ) to denote the state, action, and reward at timestep t.
A trajectory can then be described by a sequence of states, actions, and rewards: τ = ( s 0 , a 0 , r 0 , s 1 , s 2 , s 3 , . . .
, s T , a T , r T ) .
For the purposes of pretraining a Decision Transformer model, we represent trajectories using returns-to-go ˆ R t =
(cid:80)
Tt ′
=
t r ′ t instead of rewards directly, resulting in the following representation for training and generation: τ = (cid:16) ˆ R 1 , s 1 , a 1 , ˆ R 2 , s 2 , a 2 , . .
, ˆ R T , s T , a T (cid:17) .
Conditional Latent Planners (CLaP)
In this section, we describe the training and inference procedures for CLaP.
First, we discuss how
CLaP is pretrained; this builds off of the Decision Transformer’s training phase (Chen et al., 2021).
We then discuss how CLaP’s encoder and decoder are trained.
The encoder maps states into latents, from which the decoder infers what actions to take.
Finally, we explain how CLaP is used at inference time.
4.1 Pretraining
CLaP builds on the Decision Transformer architecture.
We denote the policy as π base θ .
We pretrain using the same objective as the Decision Transformer, namely to minimize error between the groundtruth action from the dataset a t and the predicted action from the sequence-based policy ˆ a t over a context length of C .
[Equation]
[Figure]
The above mean-squared error is used for environments with continuous action spaces.
In the discrete setting, the loss can be swapped with cross-entropy loss (Chen et al., 2021).
4.2 Encoder Training
The encoder maps states to low-dimensional latents, f θ : S → Z .
We model latent structure as in
Wang et al.
(2021), where the latents follow a Brownian bridge process.
The Brownian bridge process has the following density pinned at a start point z 0 at t = 0 and end point z T at t = T :
[Equation]
This density function denotes how z t should be more like z 0 at the start, z T at the end, and linearly interpolated in between but with greater uncertainty. Figure 1 illustrates this uncertainty in the
Brownian bridge process.
We use the same contrastive objective as in Wang et al.
The objective maps a triplet of states s 1 , s 2 , s 3 to f θ ( s 1 ) , f θ ( s 2 ) , f θ ( s 3 ) and fits a Brownian bridge density in Equation 2.
Formally, given multiple sequences of agent states in trajectories, S = { s 1 , . . .
, s N } , we draw batch
B : { ( s 0 , s t , s T ) } of randomly sampled positive triplets s 0 , s t , s T where 0 < t < T .
Then, our encoder is optimized by
[Equation]
where
[Equation]
This objective can be viewed as maximizing the extent to which true triplets from agent trajectories follow the Brownian bridge process while minimizing the extent to which alternate midpoints sampled from other trajectories do so.
One crucial difference between Wang et al.
and CLaP is that Wang et al.
(2021) maps all start points to a single latent, and all end points to another latent; their start points are centered at 0 and end points at 1 .
Our work does not pin start and end latents because we want to learn a policy that can generalize to any start and goal state.
4.3 Decoder Training with Latent Paths
After training the encoder, we train the decoder to produce actions using information from latent space.
Let (cid:16) ˆ R 1 , s 1 , a 1 , ˆ R 2 , s 2 , a 2 , . . .
, ˆ R C , s C (cid:17) denote the original input sequence, where C is the [Figure] context length.
We finetune the policy to produce an action similar to ground-truth action a C .
The finetuning process is as follows.
We first edit the trajectory to incorporate future information for t ∈ [1 , C ] .
We fetch the next state s t +1 ∈ R d and encode the state with the trained encoder f θ .
We now have the latent embedding z t +1 ∈ R d ′ .
To provide the policy with information about the immediate future latent, we add s t and z t +1 together.
If the two do not share dimension size, we learn a matrix M ∈ R d × d ′ which maps the latent to the state space dimensions.
The new state is then ˆ s t =
s t + Mz t
+1 .
We finetune our pretrained model from Section 4.1 with the new input sequence ( ˆ R 1 , ˆ s 1 , a 1 , . . .
, ˆ R C , ˆ s C ) .
4.4 CLaP at inference time
Figure 2 illustrates the inference setup.
Given the initial state s 0 and goal state s T , we encode them to produce the latent codes z 0 and z T .
We sample a latent path with Brownian bridge dynamics in
Equation 2.
This produces a sequence ( z 0 , z 1 , . . .
, z T ) .
We then decode actions conditioned on the past input sequence of context-length C and the future latent target: at step t , CLaP takes in the sequence ( R t − C , ˆ s t − C , a t − C , . . .
, R t , ˆ s t ) which already incorporates z t +1 in ˆ s t , and generates a t .
When the goal state is not known, we take it to be the average of the terminal states.
Similarly, when
T is unknown, we take it to be the average sequence length from our training data.
This is a limitation of our method, particularly when dealing with completely new distributions at test time.
In realistic scenarios, we can expect the goal to be known ahead of time.
Despite these limitations, we find that
CLaP still performs well in D4RL simulated environments, suggesting the learned latent dynamics is simple enough to generalize to unseen states.
Experiments
In this section, we investigate the performance of CLaP compared to two state-of-the-art methods.
One is a non-sequential method, Conservative Q-Learning (CQL) (Kumar et al., 2020), and the other is a sequence-based method, Decision Transformer (DT) (Chen et al., 2021).
We focus on how CLaP compares to DT, analyzing how much sequence models can benefit from latent goal-conditioning and incremental decoding from a latent path.
We evaluate on two continuous control tasks from the D4RL benchmark, which require fine-grained continuous control (Fu et al., 2020).
Maze2D is a navigation task where a 2D agent tries to reach a fixed goal location.
The Adroit domain requires control of adegree-of-freedom simulated robotic hand to hammer a nail, open a door, twirl a pen, or pick up and move a ball.
For the four adroit domain tasks, we consider the human and expert data settings.
Figure 3 illustrates all the environments used in our work.
We use the default hyperparameters detailed in
Chen et al.
(2021) for training, which is done on a single Nvidia Titan Xp GPU.
[Table] We first compare sequential models DT and CLaP with the non-sequential model CQL.
CQL represents the state-of-the-art non-sequential model-free offline RL method.
CQL results are reported by the original paper (Kumar et al., 2020).
Scores are normalized between 0 and 100, representing random and expert performance, respectively.
Our results are shown in Table 1.
We find that DT
and
CLaP almost always outperform CQL, demonstrating that we generally benefit using high-capacity sequential models for decision making.
We now analyze how goal-conditioning helps sequence models by comparing the two sequence models.
CLaP achieves competitive performance in a majority of tasks and shows pronounced improvement over base DT in domains with complex transition dynamics: 15.5%, 49.4%, and 149.8% boosts in Large Maze, Pen, and Hammer Human.
This finding suggests that CLaP and its goal-conditioned latent paths effectively reduce the burden on the model from learning the complex transition dynamics of these environments.
This enables the policy to more easily decode good actions, leading to better performance.
Conclusion
We proposed Conditional Latent Planners, seeking to incorporate goal-conditioned information into sequence models for decision making.
On the D4RL benchmark, we showed CLaP matches the performance of DT, while outperforming on tasks that have particularly complex transition dynamics.
This finding suggests that CLaP and its directed diffusion effectively reduce the burden on the model from learning the complex transition dynamics of these environments.
While our work empirically showed promising preliminary results with applying goal-conditioning even in domains with limited data, it is nevertheless challenging to learn good latent structures in high-variance, low-data scenarios.
Future work can research how to learn better latent representations even in settings with limited data so that sequence models for decision making can benefit maximally from goal-conditioning. Acknowledgments
Harry D. Shin’s research was supported by the CURIS Summer Internship Program at Stanford
University.
Rose E. Wang is supported by the National Science Foundation Graduate Research
Fellowship.
The authors would like to give thanks to CoCoLab principal investigator Noah Goodman and the rest of CoCoLab for their support and helpful discussions.