LAD: Language Augmented Diffusion for Reinforcement Learning
Abstract
Learning skills from language provides a powerful avenue for generalization in reinforcement learning, although it remains a challenging task as it requires agents to capture the complex interdependencies between language, actions, and states.
In this paper, we propose leveraging L anguage A ugmented D iffusion models as a planner conditioned on language (LAD).
We demonstrate the comparable performance of LAD with the state-of-the-art on the CALVIN language robotics benchmark with a much simpler architecture that contains no inductive biases specialized to robotics, achieving an average success rate (SR) of 72% compared to the best performance of 76%.
We also conduct an analysis on the properties of language conditioned diffusion in reinforcement learning.
Introduction
It has been a longstanding dream of the AI community to be able to create a household robot that can follow natural language instructions and execute behaviors such as cleaning dishes or organizing the living room [1, 2, 3, 4, 5, 6, 7, 8, 9].
Incorporating language into Reinforcement Learning (RL) has great potential for generalization, enabling agents to utilize common sense priors across tasks and environments.
Language provides an expressive abstraction of the environment and systematic generalization to new actions[10, 11, 12, 13, 14].
Given the recent progress in Natural Language
Processing, how can one incorporate the powerful capabilities of language models and utilize them for downstream decision-making?
Language-conditioned policies (LCPs) are one class of policies in RL [15, 16, 13] that can be used to formulate this task, through the conditioning of behavior on natural language instruction.
In this paper, we consider a novel approach for constructing LCPs by viewing the image as a sequence of state-actions rather than pixels, reformulating language conditioned decision-making into text-toimage generation.
Diffusion models such as DALL-E 2 [17] and GLIDE [18] have recently shown promise as generative models, with state-of-the-art text-to-image generation results demonstrating a surprisingly deep understanding of semantic relationships and generation of novel scenes.
A key driver of the recent success in generative modeling is the usage of classiﬁer-free guidance, which is amenable to the RL framework through the usage of language as a reward function.
Inspired by the recent success of diffusion models in generative modeling [17, 18], we propose a new algorithm for LCPs (LAD) via latent diffusion models [19].
We demonstrate comparable performance to the state of the art on the
CALVIN benchmark [20], with average success rate (SR) of 72% compared to the best performance of 76%.
Work done during a summer internship at Meta AI 2022.
Language and Reinforcement Learning Workshop, 36th Conference on Neural Information Processing Systems (NeurIPS 2022).
Language Augmented Diffusion 2.1 Method Overview
We consider a language-conditioned RL setting where we assume that the true reward function R is unknown, and must be inferred from a natural language instruction L ∈ L .
Formally, let F be the function space of R .
Then the goal becomes learning a φ : L (cid:55)→ F , and maximizing the policy objective conditioned on the reward function φ ( L ) , J ( π ( · | s, R ))
=
E a ∼ π,s ∼ p (cid:80) ∞ t =
γ t r t .
Note that the space of tasks that can be speciﬁed by language is much larger than that of reward, due to the Markov restriction of the latter.
For example, “pour the milk” and “pour the milk after ﬁve o’clock” are both valid instructions, but are indistinguishable from a reward function if the state does not contain temporal information.
We assume access to a prior collected dataset D of N annotated trajectories τ i = (cid:104) ( s 0 , a 0 , ...
s T ) , L i (cid:105) .
The language conditioned policy π β , or the behavior policy.
is deﬁned to be the policy that generates the aforementioned dataset.
In this paper, we assume access to a dataset of expert trajectories, such that π β = optimal policy π (cid:63) .
In this case, the policy objective reduces to imitation learning, or
[Equation]
As we tackle the problem from a planning perspective, we deﬁne a trajectory generator as P and switch the atomic object from actions to trajectories τ .
Thus we aim to
[Equation]
To model this, we turn to diffusion models [21].
Inspired by non-equilibrium thermodynamics, the common forms of diffusion models [22, 23, 24] propose modeling the data distribution p ( τ ) as a random process that steadily adds increasingly varied amounts of Gaussian noise to samples from p ( τ ) until the distribution converges to the standard normal.
We denote the forward process as f ( τ t | τ t − 1 ) , with a sequence of variances ( β 0 , β 1 ...β T ) .
We deﬁne α t := 1 − β t and ¯ α t :=
(cid:81) ts [-=
-]=1 α s .
f ( τ 1: T | τ 0 ) =
T (cid:89) t = f
( τ t | τ t − 1 ) , where f ( τ t | τ t − 1 ) =
N ( τ t ; (cid:112) 1 − β t τ t − 1 , β t I ) .
One can tractably reverse this process when conditioned on τ 0 , which allows for the construction of a sum of the typical variational lower bounds for learning the backward process’ density function [22].
Since the backwards density also follows a Gaussian, it sufﬁces to predict µ θ and Σ θ which parameterize the backwards distribution:
[Equation]
In practice, Σ θ is often ﬁxed to constants, but can also be learned through reparameterization.
Following [23] we consider learning only µ θ , which can be computed just as a function of τ t and (cid:15) θ ( τ t , t ) .
One can derive that τ t = √ ¯ α t τ 0 + √ 1 − ¯ α t (cid:15) for (cid:15) ∼ N ( 0 , I ) , through a successive reparameterization of (3) until arriving at f ( τ t | τ 0 ) .
To sample from p ( τ ) , we need only to learn (cid:15) θ , which is done by regressing to the ground truth (cid:15) given by the tractable backwards density.
Assuming we have (cid:15) θ , we can then follow a Markov chain of updates that eventually converges to the original data distribution, in a procedure reminiscent of Stochastic Gradient Langevin Dynamics [25]:
[Equation]
where z ∼ N ( 0 , I ) [Equation]
Thus, by using a variant of (cid:15) θ conditioned on language to denoise our latent plans, we can effectively model −∇ τ P β ( τ | R ) with our diffusion model, iteratively guiding our generated trajectory towards the optimal trajectories conditioned on language.
2.2 Model Architecture
It is computationally infeasible to operate directly on the pixel space.
Instead, we do planning in latent state space by ﬁrst compressing the visual input with a β -TCVAE [26], similar to latent diffusion models [19].
We choose β -TCVAE for its efﬁciency in compression.
The disentangled representation is a property that is also worth taking into consideration, and its effects on the denoising training process.
It is still unclear whether disentanglement is beneﬁcial or not.
We adopt
CLIP [27] as our textual encoder.
CLIP is trained on large-scale image-text pairs and is able to align visual and textual input in its embedding space.
Speciﬁcally, we use the Transformer variant [28] as the text encoder with the architecture modiﬁcations described in [29].
We use a modiﬁed temporal
U-Net[30], which performs 1D convolution only across the time dimension, rather than the 2D convolution typical in text-to-image generation.
This is motivated by our wish to preserve equivariance along the time dimension but not the state-action dimension.
We modify the architecture in [30] by adding conditioning via cross attention in a fashion that resembles the latent diffusion model, but also uses a temporal convolution for projection to the token embeddings rather than the traditional 2D convolution.
We use DDIM [24] during inference for increased computational efﬁciency and faster planning.
DDIM uses strided sampling and is able to capture almost the same level of ﬁdelity as typical DDPM sampling [23] with an order of magnitude speedup.
For rolling out the latent plans generated by the denoiser, we resample a new plan with the frequency of H , until either the task is completed successfully or the maximum timestep is reached.
In between samples, we roll out the open-loop plan without replanning.
We set the total time horizon equal to 3 H in our experiments, which means at maximum we will sample three latent plans from the model.
Experiments 3.1 Dataset and Metric
We evaluate a subset of the CALVIN benchmark [20] due to time, selecting the ﬁve most common subtasks (i.e tasks with the most pre-collected data) that do not belong to the same subtask category.
This metric allows for fast experimentation since there is signiﬁcantly less data to converge on, while still covering a diverse and broad range of tasks as we enforce that each task will belong to a different category.
After pretraining our VAE on all data, we freeze the autoencoder and train the
We pass sampled noise to our denoising autoencoder along with an encoded language prompt.
The autoencoder is instantiated as a temporal U-Net.
By repeating this process iteratively, we are able to generate high-ﬁdelity latent plans conditioned on language.
We are also able to decode the latent states into pixel space to analyze and interpret the plans generated by the denoising autoencoder.
U-Net just with data from the ﬁve tasks.
This setup is motivated by the assumption that we only wish to evaluate the performance of the downstream decision-making, so we assume that there already exists a robust pre-trained autoencoder.
In order to provide a fair comparison, we obtain results from our comparisons by freezing the autoencoders of their ﬁnal checkpoint which was trained on all data and training the rest of the model on just the ﬁve tasks.
We roll out all evaluated policies for 5 trajectories per task, for a total of 25 rollouts per policy.
We compare the two strongest models released on this benchmark so far, HULC and MCIL [9, 31] All comparisons are trained in their ofﬁcial repository 1 .
[Figure]
3.2 Effects of Diffusion
An interesting phenomenon we observe when rolling out with the diffusion model is its ability to robustly model arbitrary starting points of trajectories, leading to an ability to recover from failures in rollouts and attempt tasks again.
For this reason, we ﬁnd that simply by replanning more times instead of just rolling out a single plan improves the performance of LAD.
There is much more to explore in this direction, and future rollout strategies that include conditioning on more than the last state or replanning more often than just once every H timesteps will likely lead to substantial gains in performance.
Conclusion
Learning the atomic sub-skills is critical to solving the multi-task planning problem and enabling the solving of more complex and open environments through state and temporal abstraction.
We explore reformulating the language-conditioned planning process as the text-to-plan generation to better learn the alignment between language and state-action pairs.
Experiments and qualitative analysis demonstrate the simplicity and effectiveness of our model.
Future work looks to extend to the long horizon setting and further probing of the generalization and compositional capabilities of the model through classiﬁer-free guidance or improved sampling methods, as well as incorporating value functions for guiding towards optimality.
Acknowledgments
We would like to thank Ashley Jiang for creating the main ﬁgures, Michael Janner and Weixi Feng for fruitful discussions, and Oier Mees and Luka Shermann for clariﬁcations and support on the
CALVIN benchmark.
B Background and Related Work Reinforcement Learning
We formulate the RL framework as Markov decision process (MDP) M
= ( S , A , R , γ, p ), with state space S , action space A , reward function R , discount factor γ , and transition dynamics p .
At each time step t , agents observe a state s ∈ S ⊆ R n , take an action a ∈ A ⊆ R m , and transition to a new state s (cid:48) with reward r following s (cid:48) , r ∼ p ( · , ·| s, a ) .
The goal of RL is then to learn either a deterministic policy π : S (cid:55)→ A or a stochastic policy where a ∼ π ( ·| s ) with a ∈ A that maximises the policy objective, J ( π ) =
E a ∼ π,s (cid:48) ∼ p (cid:80) ∞ t =
γ t r t .
Notice this is nothing but the expected discounted cumulative reward or expected return.
Language-Conditioned Policy Language-conditioned policies [15, 16, 13] have been explored in the reinforcement learning community to improve the ability to abstract the goal and the generalization to a new environment [10, 11, 12, 13].
However, these LCPs still struggle with long-horizon language commands.
We are the ﬁrst to leverage the advantage of the diffusion model in compositionality and long-horizon decision-making to address such challenges.
Diffusion in Ofﬂine Reinforcement Learning
Given the success of the denoising diffusion probabilistic models [23] (DPM) applied in text-toimage synthesis [32], the DPM has been further explored in both discrete and continuous data domains, including image and video synthesis [33, 34], text generation [35], and time series [36].
Diffusion planning [30] ﬁrst proposed to transform the planning problem into inpainting and utilize diffusion models to solve the problem.
Speciﬁcally, they diffuse the state and actions jointly to implement imitation learning and goal-conditioned reinforcement learning.
This leverage the diffusion to solve long-horizon and compositionality issue in planning.
Instead of predicting the whole trajectory for each state, [37] apply the diffusion model to sample a single action at a time conditioned with states.
Analogous to the Diffuser [30] in trajectory-planning perspective and [37] in ofﬂine policy-optimization perspective, we are in the language-grounding perspective.