OPAL: O FFLINE P RIMITIVE D ISCOVERY FOR A CCELERATING O FFLINE R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Reinforcement learning (RL) has achieved impressive performance in a variety of online settings in which an agent’s ability to query the environment for transitions and rewards is effectively unlimited.
However, in many practical applications the situation is reversed: an agent may have access to large amounts of undirected ofﬂine experience data, while access to the online environment is severely limited.
In this work, we focus on this ofﬂine setting.
Our main insight is that, when presented with ofﬂine data composed of a variety of behaviors, an effective way to leverage this data is to extract a continuous space of recurring and temporally extended primitive behaviors before using these primitives for downstream task learning.
Primitives extracted in this way serve two purposes: they delineate the behaviors that are supported by the data from those that are not, making them useful for avoiding distributional shift in ofﬂine RL; and they provide a degree of temporal abstraction, which reduces the effective horizon yielding better learning in theory, and improved ofﬂine RL in practice.
In addition to beneﬁting ofﬂine policy optimization, we show that performing ofﬂine primitive learning in this way can also be leveraged for improving few-shot imitation learning as well as exploration and transfer in online RL on a variety of benchmark domains.
Visualizations are available at https://sites.google.com/view/opal-iclr
I NTRODUCTION
Reinforcement Learning (RL) systems have achieved impressive performance in a variety of online settings such as games (Silver et al., 2016; Tesauro, 1995; Brown & Sandholm, 2019) and robotics (Levine et al., 2016; Dasari et al., 2019; Peters et al., 2010;
Parmas et al., 2019; Pinto &
Gupta, 2016), where the agent can act in the environment and sample as many transitions and rewards as needed.
However, in many practical applications the agent’s ability to continuously act in the environment may be severely limited due to practical concerns (Dulac-Arnold et al., 2019).
For example, a robot learning through trial and error in the real world requires costly human supervision, safety checks, and resets (Atkeson et al., 2015), rendering many standard online RL algorithms inapplicable (Matsushima et al., 2020).
However, in such settings we might instead have access to large amounts of previously logged data, which could be logged from a baseline hand-engineered policy or even from other related tasks.
For example, in self-driving applications, one may have access to large amounts of human driving behavior; in robotic applications, one might have data of either humans or robots performing similar tasks.
While these ofﬂine datasets are often undirected (generic human driving data on various routes in various cities may not be directly relevant to navigation of a speciﬁc route within a speciﬁc city) and unlabelled (generic human driving data is often not labelled with the human’s intended route or destination), this data is still useful in that it can inform the algorithm about what is possible to do in the real world, without the need for active exploration.
In this paper, we study how, in this ofﬂine setting, an effective strategy to leveraging unlabeled and undirected past data is to utilize unsupervised learning to extract potentially useful and temporally extended primitive skills to learn what types of behaviors are possible .
While this might seem intuitive at ﬁrst, it is actually a bit counterintuitive, given the fact that prior works that performed controlled empirical analysis on the beneﬁts of hierarchical RL (Nachum et al., 2019b) ﬁnd that hierarchies work better than standard online RL because of improved exploration and there is no exploration to be done in an ofﬂine RL setting. We hypothesize that these primitives provide us
[Figure]
with a compact space to learn our task policy on, implicitly constrain the task policy learned on compact space to stay close to the ofﬂine data distribution, and reduce the effective time horizon of the task, thereby leading to improved performance in ofﬂine RL by reducing the curse of error propagation.
For example, consider a dataset of an agent performing undirected navigation in a maze environment (Figure 1) While the dataset does not provide demonstrations of exclusively one speciﬁc point-to-point navigation task, it nevertheless presents clear indications of which temporally extended behaviors are useful and natural in this environment (e.g., moving forward, left, right, and backward), and our unsupervised learning objective aims to distill these behaviors into temporally extended primitives.
Once these locomotive primitive behaviors are extracted, we can use them as a compact constrained temporally-extended action space for learning a task policy with ofﬂine RL, which only needs to focus on task relevant navigation, thereby making task learning easier.
For example, once a speciﬁc point-to-point navigation is commanded, the agent can leverage the learned primitives for locomotion and only focus on the task of navigation, as opposed to learning locomotion and navigation from scratch.
We refer to our proposed unsupervised learning method as Ofﬂine Primitives for Accelerating ofﬂine reinforcement Learning (OPAL), and apply this basic paradigm to ofﬂine RL, where the agent is given a single ofﬂine dataset to use for both the initial unsupervised learning phase and then a subsequent task-directed ofﬂine policy optimization phase.
Despite the fact that no additional data is used, we ﬁnd that our proposed unsupervised learning technique can dramatically improve ofﬂine policy optimization compared to performing ofﬂine policy optimization on the raw dataset directly.
In addition to these empirical results, we provide a theoretical foundation for our method, articulating what properties the initial ofﬂine dataset should possess in order to guarantee good or improved performance on a speciﬁc downstream ofﬂine RL task.
To the best of our knowledge, ours is the ﬁrst work to theoretically justify and experimentally verify the beneﬁts of primitive learning in ofﬂine RL settings, showing that hierarchies can provide temporal abstraction that allows us to reduce the effect of compounding errors issue in ofﬂine RL.
Beyond ofﬂine RL, and although this isn’t the main focus of the work, we also show the applicability of our method for accelerating RL by incorporating OPAL as a preprocessing step to standard online RL, few-shot imitation learning, and multi-task transfer learning.
In all settings, we demonstrate that the use of OPAL can improve the speed and quality of downstream task learning.
R ELATED W ORK
Ofﬂine RL.
Ofﬂine RL presents the problem of learning a policy from a ﬁxed prior dataset of transitions and rewards.
Recent works in ofﬂine RL (Kumar et al., 2019; Levine et al., 2020; Wu et al., 2019; Ghasemipour et al., 2020; Jaques et al., 2019; Fujimoto et al., 2018) constrain the policy to be close to the data distribution to avoid the use of out-of-distribution actions (Kumar et al., 2019; Levine et al., 2020).
To constrain the policy, some methods use distributional penalties, as measured by KL divergence (Levine et al., 2020; Jaques et al., 2019), MMD (Kumar et al., 2019), or Wasserstein distance (Wu et al., 2019).
Other methods ﬁrst sample actions from the behavior policy and then either clip the maximum deviation from those actions (Fujimoto et al., 2018) or just use those actions (Ghasemipour et al., 2020) during the value backup to stay within the support of the ofﬂine data.
In contrast to these works, OPAL uses an ofﬂine dataset for unsupervised learning of a continuous space of primitives.
The use of these primitives for downstream tasks implicitly constrains a learned primitive-directing policy to stay close to the ofﬂine data distribution.
As we demonstrate in our experiments, the use of OPAL in conjunction with an off-the-shelf ofﬂine RL algorithm in this way can yield signiﬁcant improvement compared to applying ofﬂine RL to the dataset directly.
Online skill discovery.
There are a number of recent works (Eysenbach et al., 2018; Nachum et al., 2018a; Sharma et al., 2019) which use unsupervised objectives to discover skills and use the discovered skills for planning (Sharma et al., 2019), few-shot imitation learning, or online RL (Eysenbach et al., 2018; Nachum et al., 2018a).
However, these works focus on online settings and assume access to the environment.
In contrast, OPAL focuses on settings where a large dataset of diverse behaviors is provided but access to the environment is restricted.
It leverages these static ofﬂine datasets to discover primitive skills with better state coverage and avoids the exploration issue of learning primitives from scratch.
Hierarchical policy learning.
Hierarchical policy learning involves learning a hierarchy of policies where a low-level policy acts as primitive skills and a high-level policy directs the low-level policy to solve a task.
While some works (Bacon et al., 2017; Stolle & Precup, 2002; Peng et al., 2019) learn a discrete set of lower-level policies, each behaving as a primitive skill, other works (Vezhnevets et al., 2017; Nachum et al., 2018b; 2019a; Hausman et al., 2018) learn a continuous space of primitive skills representing the lower-level policy.
These methods have mostly been applied in online settings.
However, there have been some recent variants of the above works (Lynch et al., 2020;
Shankar & Gupta, 2020; Krishnan et al., 2017; Merel et al., 2018) which extract skills from a prior dataset and using it for either performing tasks directly (Lynch et al., 2020) or learning downstream tasks (Shankar & Gupta, 2020; Krishnan et al., 2017; Merel et al., 2018) with online RL.
While OPAL is related to these works, we mainly focus on leveraging the learned primitives for asymptotically improving the performance of ofﬂine RL; i.e., both the primitive learning and the downstream task must be solved using a single static dataset.
Furthermore, we provide performance bounds for OPAL and enumerate the speciﬁc properties an ofﬂine dataset should possess to guarantee improved downstream task learning, while such theoretical guarantees are largely absent from existing work.
P RELIMINARIES
We consider the standard Markov decision process (MDP) setting (Puterman, 1994), speciﬁed by a tuple M = ( S , A , P , µ, r, γ ) where S represents the state space, A represents the action space,
P ( s (cid:48) | s, a ) represents the transition probability, µ ( s ) represents the initial state distribution, r ( s, a ) ∈ ( − R max , R max ) represents the reward function, and γ ∈ (0 , 1) represents the discount factor.
A policy π in this MDP corresponds to a function S → ∆( A ) , where ∆( A ) is the simplex over A .
It induces a discounted future state distribution d π , deﬁned by d π ( s ) =
( − γ ) (cid:80) ∞ t =
γ t P ( s t = s | π ) , where P ( s t = s | π ) is the probability of reaching the state s at time t by running π on M .
For a positive integer k , we use d πk ( s ) =
( − γ k ) (cid:80) ∞ t =
γ tk P ( s tk = s | π ) to denote the every- k -step state distribution of π .
The return of policy π in MDP
M is deﬁned as J RL ( π, M ) =
11 − γ E s ∼ d π ,a ∼ π ( a | s )
[ r ( s, a )] .
We represent the reward- and discount-agnostic environment as a tuple
E = ( S , A , P , µ ) .
We aim to use a large, unlabeled, and undirected experience dataset D := { τ ri :=
( s t , a t ) c − 1 t =0 } Ni =1 associated with
E to extract primitives and improve ofﬂine RL for downstream task learning.
To capture multi-modality in D , we assume it is generated by ﬁrst sampling a behavior policy π ∼ Π along with an initial state s ∼ κ , where Π ,
κ represent some (unknown) distributions over policies and states, respectively, and then running π on E for c time steps starting at s
= s .
We deﬁne the probability of a sub-trajectory τ :=
( s 0 , a 0 , . . .
, s c − 1 , a c − 1 ) in D under a policy π as π ( τ ) =
κ ( s 0 ) (cid:81)
[Equation]
t =
π ( a t | s t ) , and the conditional probability as π ( τ | s )
=
[ s = s 0 ]
[-c − 1 t =
-]
[Equation]
π ( a t | s t ) .
In this work, we will show how to apply unsupervised learning techniques to D to extract a continuous space of primitives π θ ( a | s, z ) , where z ∈ Z , the latent space inferred by unsupervised learning.
We intend to use the learned π θ ( a | s, z ) to asymptotically improve the performance of ofﬂine RL for downstream task learning.
For ofﬂine RL, we assume the existence of a dataset D r :=
{ τ ri :=
( s t , a t , r t ) c − 1 t =0 } Ni =1 , corresponding to the same sub-trajectories in
D labelled with MDP rewards.
Additionally, we can use the extracted primitives for other applications such as few-shot imitation learning, online RL, and online multi-task transfer learning.
We review the additional assumptions for these applications in Appendix A. [Table] < [Table]
Forward Autoencoding loss + KL constraint Fine-tuning with BC O ﬄ ine RL s z once per c steps every
[Table] > Primitive Policy a [Figure]
[Table]
[Figure]
[Table] [Figure]
O FFLINE RL WITH OPAL
In this section, we elaborate on OPAL, our proposed method for extracting primitives from D and then leveraging these primitives to learn downstream tasks with ofﬂine RL.
We begin by describing our unsupervised objective, which distills D into a continuous space of latent-conditioned and temporally-extended primitive policies π θ ( a | s, z ) .
For learning downstream tasks with ofﬂine RL, we ﬁrst label D r with appropriate latents using the OPAL encoder q φ ( z | τ ) and then learn a policy π ψ ( z | s ) which is trained to sample an appropriate primitive every c steps to optimize a speciﬁc task, using any off-the-shelf ofﬂine RL algorithm.
A graphical overview of ofﬂine RL with OPAL is shown in Figure 2.
While we mainly focus on ofﬂine RL, we brieﬂy discuss how to use the learned primitives for few-shot imitation learning, online RL, and multi-task online transfer learning in section 5 and provide more details in Appendix A.
4. E XTRACTING TEMPORALLY - EXTENDED PRIMITIVES FROM DATA
We would like to extract a continuous space of temporally-extended primitives π θ ( a | s, z ) from D which we can later use as an action space for learning downstream tasks with ofﬂine RL.
This would reduce our effective task horizon, thereby making the downstream learning easier, as well as allow the downstream policy to stay close to the ofﬂine data distribution, thereby bringing stability to the downstream learning.
We propose the following objective for learning π θ , incorporating an auto-encoding loss function with a KL constraint to encourage better generalization:
[Equation]
where ˆ E indicates empirical expectation.
The learned components of this objective may be interpreted as encoder, decoder, and prior:
Encoder: q φ ( z | τ ) encodes the trajectory τ of state-action pairs into a distribution in latent space and gives out parameters of that distribution.
In our case, we represent q φ with a bidirectional GRU which takes in τ and gives out parameters of a Gaussian distribution ( µ encz , σ encz ) .
Decoder (aka Primitive Policy): π θ ( a | s, z ) is the latent-conditioned policy.
It maximizes the conditional log-likelihood of actions in τ given the state and the latent vector.
In our implementation, we parameterize it as a feed-forward neural network which takes in current state and latent vector and gives out parameters of a Gaussian distribution for the action ( µ a , σ a ) .
Prior/Primitive Predictor: ρ ω ( z | s 0 ) tries to predict the encoded distribution of the sub-trajectory τ from its initial state.
Our implementation uses a feed-forward neural network which takes in the initial state and gives out parameters of a Gaussian distribution ( µ prz , σ prz ) .
In practice, rather than solving the constrained optimization directly, we implement the KL constraint as a penalty, weighted by an appropriately chosen coefﬁcient β .
Thus, one may interpret our unsupervised objective as using a sequential β -VAE (Higgins et al., 2016).
However, our prior is conditioned on s 0 and learned as part of the optimization because the set of primitives active in D depends on s 0 .
If β = 1 , OPAL is equivalent to a conditional VAE maximizing log probability of τ conditioned on its initial state s 0 ; see Appendix D for more details.
Despite the similarities between our proposed objective and VAEs, our presentation of OPAL as a constrained auto-encoding objective is deliberate.
As we will show in Section 4.3, our theoretical guarantees depend on a welloptimized auto-encoding loss to provide beneﬁts of using learned primitives π θ for downstream tasks.
In contrast, a VAE loss, which simply maximizes the likelihood of observed data, may not necessarily provide a beneﬁt for downstream tasks.
For example, if the data can be generated by a single stationary policy, a VAE-optimal policy π θ can simply ignore the latent z , thus producing a degenerate space of primitives.
In contrast, when the KL constraint in our objective is weak (i.e., (cid:15) KL (cid:29) 0 or β < 1 ), the auto-encoding loss is encouraged to ﬁnd a unique z for distinct τ to optimize reconstruction loss.
4. O FFLINE RL WITH PRIMITIVES FOR DOWNSTREAM TASKS
After distilling learned primitives from D in terms of an encoder q φ ( z | τ ) , a latent primitive policy (or decoder)
π θ ( a | s, z ) , and a prior ρ ω ( z | s 0 ) , OPAL then applies these learned models to improve ofﬂine RL for downstream tasks.
As shown in Figure 2, our goal is to use a dataset with reward labeled sub-trajectories D r
= { τ i :=
( s it , a it , r it )
c − 1 t =0 } Ni =1 to learn a behavior policy π that maximizes cumulative reward.
With OPAL, we use the learned primitives π θ ( a | s, z ) as low-level controllers, and then learn a high-level controller π ψ ( z | s ) .
To do so, we relabel the dataset D r in terms of temporally extended transitions using the learned encoder q φ ( z | τ ) .
Speciﬁcally, we create a dataset D r hi =
( s i 0 , z i , (cid:80) c − 1 t =
γ t r it , s ic ) } Ni =1 , where z i ∼ q φ ( ·| τ i ) .
Given D r hi , any off-the-shelf ofﬂine RL algorithm can be used to learn π ψ ( z | s ) (in our experiments we use CQL (Kumar et al., 2020)).
As a way to ensure that the c -step transitions τ i :=
( s it , a it , r it )
c − 1 t =0 remain consistent with the labelled latent action
z
i , we ﬁnetune π θ ( a | s, z ) on
D r lo
=
(( s it , a it ) c − 1 t =0 , z i ) } Ni =1 with a simple latent-conditioned behavioral cloning loss:
[Equation]
4. S UBOPTIMALITY AND PERFORMANCE BOUNDS FOR OPAL
Now, we will analyze OPAL and derive performance bounds for it in the context of ofﬂine RL, formally examining the beneﬁt of the temporal abstraction afforded by OPAL as well as studying what properties D should possess so that OPAL can improve downstream task performance.
As explained above, when applying OPAL to ofﬂine RL, we ﬁrst learn the primitives π θ ( a | s, z ) using D , and then learn a high-level task policy π ψ ( z | s ) in the space of the primitives.
Let π ψ ∗ ( z | s ) be the optimal task policy.
Thus the low-level π θ and high-level π ψ ∗ together comprise a hierarchical policy, which we denote as π θ,ψ ∗ .
To quantify the performance of policies obtained from OPAL, we deﬁne the notion of suboptimality of the learned primitives π θ ( a | s, z ) in an MDP M with an associated optimal policy π ∗ as
[Equation]
To relate SubOpt( θ ) with some notion of divergence between π ∗ and π θ,ψ ∗ , we introduce the following performance difference lemma.
Lemma 4.0.1.
If π 1 and π 2 are two policies in M , then
[Equation]
where D TV ( π 1 ( τ | s )
||
π 2 ( τ | s )) denotes the TV divergence over c -length sub-trajectories τ sampled from π 1 vs. π 2 (see section 3).
Furthermore,
[Equation]
[Equation]
The proof of the above lemma and all the following results are provided in Appendix B.
Through above lemma, we showed that the suboptimality of the learned primitives can be bounded by the total variation divergence between the optimal policy π ∗ in M and the optimal policy acting through the learned primitives π θ,ψ ∗ .
We now continue to bound the divergence between π ∗ and π θ,ψ ∗ in terms of how representative D is of π ∗ and how optimal the primitives π θ are with respect to the auto-encoding objective (equation 1).
We begin with a deﬁnition of how often an arbitrary policy appears in Π , the distribution generating D :
Deﬁnition 1.
We say a policy π in M is ζ -common in Π if E π ∼ Π ,s ∼ κ [D TV ( π ( τ | s )
||
π ( τ | s ))] ≤ ζ .
Theorem 4.1.
Let θ, φ, ω be the outputs of solving equation 1, such that J ( θ, φ, ω ) =
(cid:15) c .
Then, with high probability 1 − δ , for any π that is ζ -common in Π , there exists a distribution H over z such that for π Hθ ( τ | s ) :
= E z ∼ H
[ π θ ( τ | z, s )] ,
[Equation]
[Equation]
where H c = E
[Equation]
∼ Π ,τ ∼
[Equation]
[Equation]
[Equation]
− 1 t =0 log π ( a t | s t )] (i.e. a constant and property of D ) and S J is a positive constant incurred due to sampling error in J ( θ, φ, ω ) and depends on concentration properties of π θ ( a | s, z ) and q φ ( z | τ ) .
Corollary 4.1.1.
If the optimal policy π ∗ of M is ζ -common in Π , and (cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13) ∞ ≤ ξ , then, with high probability 1 − δ ,
[Equation]
As we can see, SubOpt( θ ) will reduce as D gets closer to π ∗ (i.e. ζ approaches 0) and better primitives are learned (i.e. (cid:15) c decreases).
While it might be tempting to increase c (i.e. the length of sub-trajectories), as much as we can, to reduce the suboptimality, it will get practically harder to control the autoencoding loss (cid:15) c with increasing c thereby leading to an increase in overall suboptimality.
This happens because it’s harder to decode longer sub-trajactories.
Till now, we have argued that there exists some near-optimal task policy π ψ ∗ if θ is sufﬁciently learned and π ∗ is sufﬁciently well-represented in D .
Now, we will show how primitive learning can improve downstream learning, by considering the beneﬁts of using OPAL with ofﬂine RL.
Building on the policy performance analysis from Kumar et al.
(2020), we now present theoretical results bounding the performance of the policy obtained when ofﬂine RL is performed with OPAL.
Theorem 4.2.
Let π ψ ∗ ( z | s ) be the policy obtained by CQL and let π ψ ∗ ,θ ( a | s ) refer to the policy when π ψ ∗ ( z | s ) is used together with π θ ( a | s, z ) .
Let π β ≡ { π ; π ∼ Π } refer to the policy generating
D r in MDP M and z ∼ π H β ( z | s ) ≡ τ ∼ π β,s
=
s , z ∼ q φ ( z | τ ) .
Then, J ( π ψ ∗ ,θ , M ) ≥ J ( π β , M ) − κ with high probability 1 − δ where
[Equation]
where D CQL is a measure of the divergence between two policies; see the appendix for a formal statement.
The precise bound along with a proof is described in Appendix B. Intuitively, this bound suggests that the worst-case deterioration over the learned policy depends on the divergence between the learned latent-space policy D CQL and the actual primitive distribution, which is controlled via any conservative ofﬂine RL algorithm (Kumar et al. (2020) in our experiments) and the size of the latent space |Z| .
Crucially, note that comparing Equation 9 to the performance bound for CQL (Equation 6 in Kumar et al.
(2020)) reveals several beneﬁts pertaining to (1) temporal abstraction – a reduction in the factor of horizon by virtue of γ c , and (2) reduction in the amount of worst-case error propagation [Table] [Figure] due to a reduced action space |Z| vs. |A| .
Thus, as evident from the above bound [Table] the total error induced due to a combination of distributional shift and sampling is signiﬁcantly reduced when OPAL is used as compared to the standard RL counterpart of this bound which is affected by the size of the entire action space for each and every timestep in the horizon.
This formalizes our intuition that OPAL does help partly mitigate distributional shift and sampling error.
One downside of using a latent space policy is that we incur unsupervised learning error while learning primitives.
However, the unsupervised learning error terms gets dominated by other error terms which pertains to ofﬂine RL.
This is also observed in practice where it’s much easier to control unsupervised loss than errors in ofﬂine RL.
E VALUATION
In this section, we will empirically show that OPAL improves learning of downstream tasks with ofﬂine RL, and then brieﬂy show the same with few-shot imitation learning, online RL, and online multi-task transfer learning.
Unless otherwise stated, we use c = 10 and dim( Z ) = 8 .
See
Appendix C for further implementation and experimental details.
Visualizations are available at https://sites.google.com/view/opal-iclr
5. O FFLINE RL WITH OPAL
Description: We use environments and datasets provided in D4RL (Fu et al., 2020).
Since the aim of our method is speciﬁcally to perform ofﬂine RL in settings where the ofﬂine data comprises varied and undirected multi-task behavior, we focus on Antmaze medium (diverse dataset), Antmaze large (diverse dataset), and Franka kitchen (mixed and partial datasets).
The Antmaze datasets involve a simulated ant robot performing undirected navigation in a maze.
The task is to use this undirected dataset to solve a speciﬁc point-to-point navigation problem, traversing the maze from one corner to the opposite corner, with only sparse 0-1 completion reward for reaching the goal.
The kitchen datasets involves a franka robot manipulating multiple objects (microwave, kettle, etc.)
either in an undirected manner (mixed dataset) or in a partially task directed manner (partial dataset).
The task is to use the datasets to arrange objects in a desired conﬁguration, with only sparse 0-1 completion reward for every object that attains the target conﬁguration.
Baseline: We use Behavior cloning (BC), BEAR (Kumar et al., 2019), EMAQ (Ghasemipour et al., 2020), and CQL (Kumar et al., 2020) as baselines.
We compare it to CQL+OPAL, which ﬁrst uses
OPAL to distill primitives from the ofﬂine dataset before applying CQL to learn a primitive-directing high-level policy. [Table] [Table]
Results: As shown in Table [Table] CQL+OPAL outperforms nearly all the baselines on antmaze (see
Figure 1 and Figure 3 for visualization) and kitchen tasks, with the exception of EMAQ having similar performance on kitchen mixed.
To ensure fair comparison with EMAQ, we use an autoregressive primitive policy.
With the exception of EMAQ on kitchen mixed, we are not aware of any existing ofﬂine RL algorithms that achieves similarly good performance on these tasks; moreover, we are not aware of any existing online RL algorithms which solve these tasks (see Table 3 for some comparisons), highlighting the beneﬁt of using ofﬂine datasets to circumvent exploration challenges.
There are two potential reasons for OPAL’s success.
First, temporally-extended primitives could make the reward propagation learning problem easier.
Second, the primitives may provide a better latent action space than the atomic actions of the environment.
To understand the relative importance of these factors, we experimented with an ablation of CQL+OPAL that uses c = 1 to remove temporal abstraction.
In this case, we ﬁnd the method’s performance to be similar to standard CQL.
This implies that the temporal abstraction provided by OPAL is one of the main contributing factors to its good performance.
This observation also agrees with our theoretical analysis.
See Appendix E for detailed discussion.
5. F EW - SHOT IMITATION LEARNING WITH OPAL
Description: Previously, we assumed that we have access to a task reward function, but only undirected data that performs other tasks.
Now, we will study the opposite case, where we are not provided with a reward function for the new task either, but instead receive a small number of taskspeciﬁc demonstrations that illustrate optimal behavior.
Simply imitating these few demonstrations is insufﬁcient to obtain a good policy, and our experiments evaluate whether OPAL can effectively incorporate the prior data to enable few-shot adaptation in this setting.
We use the Antmaze environments (diverse datasets) to evaluate our method and use an expert policy for these environments to sample n = 10 successful trajectories.
Baseline and Results: For baselines, we use Behavior cloning (BC) and the model from Wang et al.
(2017), which prescribes using a sequential VAE (SVAE) over state trajectories in conjunction with imitation learning.
As shown in Table 2, BC+OPAL clearly outperforms other baselines, showing the importance of temporal abstraction and ascertaining the quality of learned primitives.
See
Appendix A for detailed discussion.
5. O NLINE RL AND MULTI - TASK TRANSFER WITH OPAL
Description: For online RL and multi-task transfer learning, we learn a task policy in space of primitives π θ ( a | s, z ) while keeping it ﬁxed.
For multi-task transfer, the task policy also takes in the task id
and we use c = 5 and Z = 8 .
Since the primitives need to transfer to a different state distribution for multi-task transfer, it only learns the action sub-trajectory distribution and doesn’t take in the state feedback.
See Appendix A for a detailed description of models.
For online RL, we use the Antmaze environments (diverse datasets) with sparse and dense rewards for evaluating our [Table]
method. For online multi-task transfer learning, we learn primitives with expert data from pick-andplace task and then use it to learn multi-task policy for MT10 and MT50 (from metaworld (Yu et al., 2020)), containing 10 and 50 robotic manipulation tasks which needs to be solved simultaneously.
Baseline and Results: For online RL, we use HIRO (Nachum et al., 2018b), a state-of-the-art hierarchical RL method, SAC (Haarnoja et al., 2018) with Behavior cloning (BC) pre-training on D , and Discovery of Continuous Options (DDCO) (Krishnan et al., 2017) which uses D to learn a discrete set of primitives and then learns a task policy in space of those primitives with online RL (Double DQN (DDQN) (Van Hasselt et al., 2015)).
For online multi-task transfer learning, we use
PPO (Schulman et al., 2017) and SAC (Haarnoja et al., 2018) as baselines.
As shown in Table 3 and Table 4 , OPAL uses temporal abstraction to improve exploration and thus accelerate online RL and multi-task transfer learning.
See Appendix A for detailed discussion.
D ISCUSSION
We proposed Ofﬂine Primitives for Accelerating Ofﬂine RL (OPAL) as a preproccesing step for extracting recurring primitive behaviors from undirected and unlabelled dataset of diverse behaviors.
We derived theoretical statements which describe under what conditions OPAL can improve learning of downstream ofﬂine RL tasks and showed how these improvements manifest in practice, leading to signiﬁcant improvements in complex manipulation tasks.
We further showed empirical demonstrations of OPAL’s application to few-shot imitation learning, online RL, and online multi-task transfer learning.
In this work, we focused on simple auto-encoding models for representing OPAL, and an interesting avenue for future work is scaling up this basic paradigm to image-based tasks.
Appendices A O THER A PPLICATIONS OF OPAL A. F EW - SHOT IMITATION LEARNING WITH OPAL
Additional Assumption : In addition to D , we assume access to a small number of expert demonstrations D exp
=
{ τ i :=
( s t , a t ) T − 1 t =0 } ni =1 where n (cid:28) N .
How to use with OPAL?
In imitation learning, the aim is to recover an expert policy given a small number of stochastically sampled expert demonstrations D exp =
{ τ i :=
( s t , a t ) T − 1 t =0 } ni =1 .
As in the ofﬂine RL setting, we use the primitives π θ ( a | s, z ) learned by OPAL as a low-level controller and learn a high-level policy π ψ ( z | s ) .
We ﬁrst partition the expert demonstrations into sub-trajectories
D exppar = { τ i,k :=
( s k + t , a k + t ) c − 1 t =0 for k = 0 , . .
, T − c } ni
=1 of length c .
We then use the learned encoder q φ ( z | τ ) to label these sub-trajectories with latent actions z i,k ∼ q φ ( z | τ i,k ) and thus create a dataset D exphi = { ( s ik + t , z i,k ) for k = 0 , . .
, T − c } ni =1 .
We use D exphi to learn the high-level policy π ψ ( z | s ) using behavioral cloning.
As in the ofﬂine RL setting, we also ﬁnetune π θ ( a | s, z ) with latent-conditioned behavioral cloning to ensure consistency of the labelled latents.
Evaluation Description: We receive a small number of task-speciﬁc demonstrations that illustrate optimal behavior.
Simply imitating these few demonstrations is insufﬁcient to obtain a good policy, and our experiments evaluate whether OPAL can effectively incorporate the prior data to enable fewshot adaptation in this setting.
We use the Antmaze environments (diverse datasets) to evaluate our method and use an expert policy for these environments to sample n = 10 successful trajectories.
Baseline: We evaluate two baselines.
First, we test a simple behavioral cloning (BC) baseline, which trains using a max-likelihood loss on the expert data.
In order to make the comparison fair to OPAL (which uses the ofﬂine dataset D in addition to the expert dataset), we pretrain the BC agent on the undirected dataset using the same max-likelihood loss.
As a second baseline and to test the quality of OPAL-extracted primitives, we experiment with an alternative unsupervised objective from Wang et al.
(2017), which prescribes using a sequential VAE (SVAE) over state trajectories in conjunction with imitation learning.
Results: As shown in Table 2, BC+OPAL clearly outperforms other baselines, showing the importance of temporal abstraction and ascertaining the quality of learned primitives.
SVAE’s slightly worse performance suggests that decoding the state trajectory directly is more difﬁcult than simply predicting the actions, as OPAL does, and that this degrades downstream task learning.
A. O NLINE RL WITH OPAL
Additional Assumptions : We assume online access to M ; i.e., access to Monte Carlo samples of episodes from M given an arbitrary policy π .
How to use with OPAL?
To apply OPAL to standard online RL, we ﬁx the learned primitives π θ ( a | s, z ) and learn a high-level policy π ψ ( z | s ) in an online fashion using the latents z as temporally-extended actions.
Speciﬁcally, when interacting the environment, π ψ ( z | s ) chooses an appropriate primitive every c steps, and this primitive acts on the environment directly for c timesteps.
Any off-the-shelf online RL algorithm can be used to learn ψ .
In our experiments, we use
SAC (Haarnoja et al., 2018).
To ensure that π ψ ( z | s ) stays close to the data distribution and avoid generalization issues associated with the ﬁxed π θ ( a | s, z ) , we add an additional KL penalty to the reward of the form D KL ( π ψ ( z | s ) ||
ρ ω
( z | s 0 )) .
Evaluation Description: We use Antmaze medium (diverse dataset) and Antmaze large (diverse dataset) from the D4RL task suite (Fu et al., 2020) to evaluate our method.
We evaluate using both a dense distance-based reward −(cid:107) g − ant xy (cid:107) and a sparse success-based reward 1[ (cid:107) g − ant xy (cid:107) ≤ 0 .
5] (the typical default for this task), where ant xy is the 2d position of the ant in the maze.
Baseline: To solve these tasks through online RL, we need both (i) hierarchy (i.e. learning a policy on top of primitives) which improves exploration (Nachum et al., 2019b) and (ii) unlabelled (i.e. no task reward) ofﬂine dataset which allows us to bootstrap the primitives.
This informed our choice of these three baselines.
First, to test the role of D in exploration, we use HIRO (Nachum et al., 2018b), a state-of-the-art hierarchical RL method, as a baseline.
Second, to test the role of temporal abstraction, we pre-train a ﬂat policy on D using behavioral cloning (BC) and then ﬁnetune the policy on downstream tasks with SAC.
Third, to test the quality of extracted primitives for online RL, we extract a discrete set of primitives with Deep Discovery of Continuous Options (DDCO) (Krishnan et al., 2017) and use Double DQN (DDQN) (Van Hasselt et al., 2015) to learn a task policy in the space of learned discrete primitives.
Results: As shown in Table 3, SAC+OPAL outperforms all the baselines, showing (1) the importance of D in exploration, (2) the role of temporal abstraction, and (3) the good quality of learned primitives.
Except for HIRO on Antmaze large with dense rewards, all other baselines fail to make any progress at all.
In contrast, SAC+OPAL only fails to make progress on Antmaze large with sparse rewards.
A. O NLINE M ULTI - TASK TRANSFER LEARNING WITH OPAL
Additional Assumption :
We assume the existence of M additional MDPs {M i =
( S i , A , P i , r i , γ ) } Mi =1 where the action space A , and the discount factor γ are same as those of
M .
How to use it with OPAL?
In the multi-task setting, we aim to learn near-optimal behavior policies on M MDPs {M i =
( S i , A , P i , r i , γ ) } Mi =1 .
As in the previous applications of OPAL, we learn a set of high-level policies π ψ ( z | s, i ) which will direct pretrained primitives π θ ( a | s, z ) to maximize cumulative rewards.
Since the state space in the M MDPs is potentially distinct from that in the ofﬂine dataset D , we cannot transfer the state distribution and can only hope to transfer the action sub-trajectory distribution.
Therefore, during the unsupervised training phase for learning π θ , we make the encoder and the decoder blind to the states in sub-trajectory.
Speciﬁcally, the encoder becomes ( µ encz , σ encz ) =
q φ ( z t | s t , { a t + i } c − 1 i =0 ) and is represented by a bidirectional GRU.
The decoder becomes π θ ( { a t , . .
. , a t + c − 1 }| z t ) which decodes the entire action sub-trajectory from the latent vector and is represented by a GRU.
With these state-agnostic primitives in hand, we then learn a policy π ψ ( z | s, i ) using any off-the-shelf online RL method.
In our experiments, we use Proximal
Policy Optimization (PPO) (Schulman et al., 2017).
Evaluation Description: We use the Metaworld task suite (Yu et al., 2020) to evaluate our method.
The dataset D for learning primitives consists of trajectories generated by an expert policy for a goal-conditioned pick-and-place task.
The pick-and-place task is suitable for unsupervised primitive learning because it contains all the basic operations (eg: move, grasp, place) required for performing more complex manipulation tasks in the Metaworld.
Once we have learned the primitives, we learn a policy π ψ ( z | s, i ) for the MT10 and MT50 benchmarks, where MT10 and MT50 contain 10 and 50 robotic manipulation tasks, respectively, which we need to be solved simultaneously.
In these experiments we use c = 5 and dim( Z ) = 8 .
Baseline: We use SAC (Haarnoja et al., 2018) and PPO (Schulman et al., 2017) as baselines.
Results: As shown in Table 4, PPO+OPAL clearly outperforms both PPO and SAC, showing the importance of temporal abstraction in online multi-task transfer.
B P ROOF OF T HEOREMS
We will begin by proving the following lemma which bounds the sampling error incurred by
J ( φ, θ, ω ) .
Lemma B.0.1.
With high probability 1 − δ ,
[Equation]
where S J is a constant dependent on concentration properties of π θ ( a | s, z ) and q φ ( z | τ ) .
Proof.
Applying Chebyshev’s inequality to
[Equation]
we get that, with high probability 1 − δ ,
[Equation]
Therefore
[Equation]
we have
[Equation]
We present a general performance difference lemma that will help in our proof of Lemma 4.0.1.
Lemma B.0.2.
If π 1 and π 2 are two policies in M , then
[Equation]
Proof.
Following the derivations in Achiam et al.
(2017) and Nachum et al.
(2018a)
[Equation]
we express the performance of a policy π in M in terms of linear opterators:
[Equation]
where R is a vector representation of the rewards of M
[Equation]
Π π is a linear operator mapping state distributions to state-action distributions according to π , and µ is a vector representation of the initial state distribution of M .
Accordingly, we express the performance difference of π 1 , π 2 as,
[Equation]
By the triangle inequality
[Equation]
we may bound equation 18 by
R max
[Equation]
We begin by approaching the ﬁrst term inside the parentheses of equation 19.
That ﬁrst term may be expressed as
[Equation]
Now we continue to the second term inside the parentheses of equation 19.
This second term may be expressed as
[Equation]
To combine equations 23 and 25, we note that
[Equation] [Equation]
Thus, we have
[Equation]
as desired.
Lemma 4.0.1.
If π 1 and π 2 are two policies in M , then
[Equation]
where D TV ( π 1 ( τ | s )
||
π 2 ( τ | s )) denotes the TV divergence over c -length sub-trajectories τ sampled from π 1 vs. π 2 (see section 3).
Furthermore,
[Equation]
Proof.
We focus on proving equation 28, as the subsequent derivation of equation 29 is straightforward by deﬁnition of SubOpt .
To derive equation 28, we may simply consider π 1 and π 2 acting in an “every- c -steps” version of
M , where the action space is now τ and reward are accumulated over c steps using γ -discounting.
Note that in this abstracted version of M , the max reward is 1 − γ c 1 − γ R max and the MDP discount is γ c .
Plugging this into the result of Lemma B.0.2 immediately yields the desired claim.
Theorem 4.1.
Let θ, φ, ω be the outputs of solving equation 1, such that J ( θ, φ, ω ) =
(cid:15) c .
Then, with high probability 1 − δ , for any π that is ζ -common in Π , there exists a distribution H over z such that for π Hθ ( τ | s ) :
= E z ∼ H
[ π θ ( τ | z, s )] ,
[Equation] [Equation]
where H c = E π ∼ Π ,s ∼ κ,τ ∼ π [ (cid:80) c
− 1 t =0 log π ( a t | s t )]
(i.e. a constant and property of D ) and S J is a pos- itive constant incurred due to sampling error in J ( θ, φ, ω ) and depends on concentration properties of π θ ( a | s, z ) and q φ ( z | τ ) .
Proof.
We start with application of the triangle inequality:
[Equation]
[Equation]
Taking expectation with respect to π ∼ Π , s ∼ κ on both the sides, we get
[Equation]
[Equation]
The last two inequality used Jensen’s inequality.
Let H ( z ) =
E π ∼ Π ,τ ∼ π [ q φ ( z | τ )] .
Cancelling out the dynamics and using equation 14 we get,
[Equation]
Corollary 4.1.1.
If the optimal policy π ∗ of M is ζ -common in Π , and
[Equation]
ξ , then, with
[Equation]
Proof.
Expanding lemma 4.0.1 using the above assumption, we have
[Equation]
Now, we can use theorem 4.1 to prove the corollary.
Theorem 4.2.
Let π ψ ∗ ( z | s ) be the policy obtained by CQL and let π ψ ∗ ,θ ( a | s ) refer to the policy when π ψ ∗ ( z | s ) is used together with π θ ( a | s, z ) .
Let π β ≡ { π ; π ∼ Π } refer to the policy generating
D r in MDP M and z ∼ π Hβ ( z | s ) ≡ τ ∼ π β,s
=
s , z ∼ q φ ( z | τ ) .
Then, J ( π ψ ∗ ,θ , M ) ≥ J ( π β , M ) − κ with high probability 1 − δ where
[Equation]
Proof.
We assume that the variational posterior q φ ( z | τ ) obtained after learning OPAL from D is same (or nearly same) as the true posterior p ( z | τ ) .
q φ can be used to deﬁne π Hβ ( z | s ) as z ∼ π Hβ ≡ τ ∼ π β , z ∼ q φ ( z | τ ) .
This induces an MDP M H =
( S , Z , P z , r z , γ c ) where Z is the inferred latent space for choosing primitives,
P z and r z are the latent dynamics and reward function such that s t + c ∼ P z
( s t + c
| s t , z t ) ≡ s t + i +1 ∼ P ( s t
+ i +1 |
s t
+ i , a t + i ) , a t
+ i ∼ π
( a t + i | s t
+ i , z t ) ∀ i ∈ { 0 , 1 , . . .
, c − 1 } and r z ( s t , z t ) =
(cid:80) c − 1 i =
γ i r ( s t
+ i , a t + i ) , and γ c is the new discount factor effectively reducing the task horizon by a factor of c .
π ( a | s, z ) is the primitive induced by q φ and π β .
Since q φ captures the true posterior, π ( a | s, z ) is the optimal primitive you can learn and its autoencoding loss, under true expectation, is (cid:15) ∗ c =
E π ∼ Π ,τ ∼ π,z
[Equation]
φ
[Equation]
Therefore, τ ∼ π β ≡ z ∼ π Hβ , τ ∼ π ( ·|· , z ) . π β is used to collect the data D r which induces an empirical MDP.
We refer to the empirical MDP induced by D r as ˆ M = ( S , A , ˆ P , ˆ µ, r, γ ) where
[Equation]
[-= ˆ s (cid:48) ] (cid:80) ( s,a ) ∼D 
 [ s =ˆ s,a =ˆ a ]-] and ˆ µ (ˆ s ) = (cid:80) s
∼D 1[ s
=ˆ s ] N .
We use q φ to get D r hi from
D r which induces another empirical MDP ˆ M H .
Using these deﬁnitions, we will try to bound | J ( π ψ ∗ ,θ , M ) − J ( π β , gM ) | .
Let’s break | J ( π ψ ∗ ,θ , M ) − J ( π β , M ) | into [Equation]
[Equation]
Since q φ captures the true variational posterior, τ ∼ π β ≡ z ∼ π Hβ , τ ∼ π ( ·|· , z ) and therefore, | J ( π H β , M H ) − J ( π β , M ) | = 0 .
For bounding, | J ( π ψ ∗ , M H − J ( π H β , M H ) | , we use theorem 3.
[Equation]
Now
[Equation]
we will try to bound | J ( π ψ ∗ ,θ , M ) − J ( π ψ ∗ , M H ) | .
The only difference between the two is that the primitive π θ ( a | s, z ) is used in M and the primitive π ( a | s, z ) is used in M H .
Therefore
[Equation]
we can write the above bound as
[Equation]
| .
Let’s ﬁrst bound their value function at a particular state s .
Using Lemma
[Equation] [Equation] [Equation]
we get
[Equation]
Now, we will try to bound D KL ( π ψ ∗ ,π ( ·|· ,z ) ( τ | s )
||
π ψ ∗ ,θ ( τ | s )) .
We have
[Equation]
The last equation comes from above deﬁnition of (cid:15) ∗ c and equation 14.
We will now try to bound (cid:12) (cid:12)(cid:12)
π ψ ∗ ( z | s ) π Hβ
( z | s ) (cid:12) (cid:12)(cid:12) ∞ using D CQL ( π ψ ∗ , π Hβ )( s ) .
Using deﬁnition of D CQL ( π ψ ∗ , π Hβ )( s ) , we have
D CQL
[Equation] [Equation]
⇒
[Equation]
∞
[Equation]
56) where ¯ z = arg max z
[Equation]
Therefore,
[Equation]
Using this to bound the returns, we get
[Equation]
[Equation]
[Equation]
We get κ = κ 1 + κ 2 .
We apply O to get the notation in the theorem.
C E XPERIMENT D ETAILS C. OPAL EXPERIMENT DETAILS
Encoder The encoder q φ ( z | τ ) takes in state-action trajectory τ of length c .
It ﬁrst passes the individual states through a fully connected network with 2 hidden layers of size H and ReLU activation.
Then it concatenates the proccessed states with actions and passes it through a bidirectional GRU with hidden dimension of H and 4 GRU layers.
It projects the output of GRU to mean and log standard deviation of the latent vector through linear layers.
Prior The prior ρ ω ( z | s ) takes in the current state s and passes it through a fully connected network with 2 hidden layers of size H and ReLU activation.
It then projects the output of the hidden layers to mean and log standard deviation of the latent vector through linear layers.
Primitive Policy
The primitive (i.e. decoder)
π θ ( a | s, z ) has same architecture as the Prior but it takes in state and latent vector and produces mean and log standard deviation of the action.
For kitchen environments, we use an autoregressive primitive policy with same architecture as used by
EMAQ (Ghasemipour et al., 2020).
We use H = 200 for antmaze environments and H = 256 for kitchen environments.
In both cases,
OPAL was trained for 100 epochs with a ﬁxed learning rate of 1 e − 3 , β = 0 . 1 (Lynch et al., 2020), Adam optimizer (Kingma & Ba, 2014) and a batch size of 50 .
C. T ASK P OLICY ARCHITECTURE
In all environments, for task policy, we used a fully connected network with 3 hidden layers of size 256 and ReLU activation.
It then projects the output of the hidden layers to mean and log standard deviation of the latent vector through linear layers.
C. SAC H YPERPARAMETERS
We used SAC (Haarnoja et al., 2018) for online RL experiments in learning a task policy either in action space A or latent space Z .
For the discrete primitives extracted from DDCO (Krishnan et al., 2017), we used Double DQN Van Hasselt et al.
We used the standard hyperparameters for SAC and Double DQN as provided in rlkit code base ( https://github.com/vitchyr/ rlkit ) with both policy learning rate and q value learning rate as 3 e − 4 .
C. CQL H YPERPARAMETERS
We used CQL (Kumar et al., 2020) for ofﬂine RL experiments in learning a task policy either in action space A or latent space Z .
We used the standard hyperparameters, as mentioned in Kumar et al.
, with minor differences.
We used policy learning rate of 3 e − 5 , q value learning rate of 3 e − 4 , and primitive learning rate of 3 e − 4 .
For antmaze tasks, we used CQL( H ) variant with τ = 5 and learned α .
For kitchen tasks we used CQL( ρ ) variant with ﬁxed α = 10 .
In both cases, we ensured α never dropped below 0 . [Table]
D C ONNECTION BETWEEN OPAL AND VAE OBJECTIVES
We are given an undirected, unlabelled and multi-modal dataset D of sub-trajectories of length c .
We would like to ﬁt a sequential VAE model to D which maximizes
[Equation]
where s 0 is initial state of the sub-trajectory.
Let’s consider
[Equation]
(using Jensen’s inequality)
[Equation]
Using the above equation
[Equation]
we have the following lower-bound for our objective function
[Equation]
We separate the parameters of decoder from prior and hence write p θ ( z | s 0 )
=
ρ ω ( z | s 0 ) .
We can expand log p θ ( τ | z, s 0 ) =
[Equation]
| s t , z ) . Since P is ﬁxed it can be removed from the objective function.
Therefore, we can write the objective function as
[Equation]
where β
= 1 .
This is similar to the autoencoding loss function we described in section 4. E A BLATION S TUDIES
As shown in Table 5, we experimented with different choices of dim( Z ) on antmaze-medium (diverse).
Using the hyperparameters from Nachum et al.
(2018a)
, we ﬁxed c = 10 .
While dim( Z ) = 8 , 16 gave similar performances, dim( Z ) = 4 performed slightly worse.
Therefore, we selected dim( Z ) =
for our ﬁnal model as it was simpler.
Temporal abstraction actually helps: To empirically verify that the gain in performance was due to temporal abstraction and not better action space learned through latent space, we tried c =( dim( Z ) = 8 ) and found the performance to be similar to that of CQL (i.e. 55 .
3 ± 3 . 8 ) thereby empirically supporting the theoretical beneﬁts of temporal abstraction.
We found dim( Z ) = 8 and c = 10 to work well with other environments as well.
However, we acknowledge that the performance of CQL+OPAL can be further improved by carefully choosing better hyperparameters for each environment or by using other ofﬂine hyperparameter selection methods for ofﬂine RL, which is a subject of future work.
F A LTERNATIVE METHODS FOR EXTRACTING PRIMITIVES FROM OFFLINE DATA
We describe alternative methods for extracting a primitive policy from ofﬂine data.
These methods are ofﬂine variants of CARML (Jabri et al., 2019) and DADS (Sharma et al., 2019).
We tried these
[Table]
[Table]
techniques in an early phase of our project and used the environment antmaze-medium (diverse) to evaluate these methods.
Let’s consider an ofﬂine multi-modal dataset D
= { ( s it , a it )
c − 1 t =0 } Ni =1 .
Let τ =
( s t ) c − 1 t =0 represent the state trajectory.
To extract primitives, we ﬁrst cluster the trajectories by maximizing the mutual information between the state trajectory τ and latent variable z (indicating cluster index) with respect to the parameters of joint distribution p φ,ω ( τ, z ) =
p ω ( z ) p φ ( τ | z ) .
For now, we consider p ω ( z ) =
Cat( p 1 , . . .
, p k ) (i.e. discrete latent variables sampled from a Categorical distribution) and represent z as one-hot vector of dimension k .
The choice of p ω ( z ) is consistent with the choices made in Jabri et al.
and Sharma et al.
Since z is discrete, we can use Bayes rule to calculate p φ,ω ( z | τ ) as
[Equation]
Our objective function becomes
[Equation]
Ofﬂine CARML and ofﬂine DADS differ only in how they model p φ ( τ | z ) :
• Ofﬂine CARML We model p φ ( τ | z ) =
[Equation]
c − 1 t =0 p φ ( s t | z ) and hence, log p φ ( τ | z ) = (cid:80) c − 1 t =0 log p φ ( s t | z ) .
• Ofﬂine DADS
We model p φ ( τ | z ) =
p ( s 0 )
[Equation]
1 t =1 p φ ( s t | s t − 1 , z ) and hence, log p φ ( τ | z ) = log p ( s 0 )
+
[Equation]
0 log p φ ( s t | s t − 1 , z ) .
Here, we only model p φ ( s t | s t − 1 , z ) and not p ( s 0 ) .
Since, log is additive in nature, p ( s 0 ) will be ignored while calculating gradient.
To optimize equation 68, we use Algorithm 2 from Jabri et al.
Once we have clustered the state trajectories τ with labels z by maximizing I ( τ ; z ) , we can use behavioral cloning (BC) to learn π θ ( a | s, z ) .
Finally, we use p φ,ω ( z | τ ) to label the reward-labelled data D r = { ( s it , a it , r it )
c − 1 t =0 } Ni =1 with latents, and transform it into D r
{+hi = { ( s i 0 , z i , (cid:80) c − 1 t =
+}
γ i r it , s c ) Ni =1 } .
The task policy π ψ is trained on D r hi using Conservative Q Learning (CQL) (Kumar et al., 2020).
Since the primitive policy π θ is trained after p φ,ω
( z | τ ) is fully trained, it doesn’t need any additional ﬁnetuning.
F. R ESULTS
Using the hyperparameters from Nachum et al.
(2018a)
, we used c = 10 .
We experimented with different values of k = 5 , 10 , 20 and found that k = 10 , 20 works the best (see Table 6 for more details).
We went with k = 10 as our ﬁnal model since it’s simpler.
Ofﬂine CARML effectively uses only 6 skills as the other 4 skills had p ω ( z )
= 0 .
Ofﬂine DADS uses all the skills.
The results are described in Table 7.
The resulting trajectory clusters (using a subset of the dataset) from discrete skills are also visualized in Figure 4 where different colors represent different clusters.
[Figure]
Since ofﬂine CARML treats the states in the trajectory conditionally independent of each other given z , the clustering mainly focuses on the spatial location.
Therefore, ofﬂine CARML isn’t able to separate out different control modes starting around the same spatial locations which explains its poor performance when combined with CQL.
On the other hand, ofﬂine DADS treats the state transitions in the trajectory conditionally independent of each other given z and thus clusters trajectories with similar state transitions together.
This allows it to more effectively separate out the control modes.
Therefore, CQL+ofﬂine DADS slightly improves upon CQL but is still limited by discrete number of skills.
Furthermore, increasing the number of skills from 10 to 20 gives similar performance.
Moreover, in these methods, it’s intractable to use continuous skill space since we use Bayes rule to calculate p φ,ω
( z | τ ) .
Therefore, we decided to switch to learning a β -VAE (Higgins et al., 2016) style generative model with continuous skill space i.e. OPAL.
F. T RAINING D ETAILS
For clustering by optimizing equation 68, both ofﬂine CARML and ofﬂine DADS only considers the global x-y pose of the ant and ignores other dimensions of the state space.
These methods fail to work when considering the full state space.
Ofﬂine CARML p φ ( s | z ) takes in the latent one-hot vector z and passes it through a fully connected network with 2 hidden layers of size H = 200 and ReLU activation.
It then projects the output of the hidden layers to mean and log standard deviation of the reduced state s (only global x-y pose considered) through linear layers.
Ofﬂine DADS p φ ( s t | s t − 1 , z ) has the same architecture as p φ ( s | z ) but also takes in the reduced state from the previous timestep.
Primitive Policy
The primitive policy π θ ( a | s, z ) takes in the current state s and latent one-hot vector z and passes it through a fully connected network with 2 hidden layers of size H = 200 and ReLU activation.
It then projects the output of the hidden layers to mean and log standard deviation of action through linear layers.
We perform the clustering for 25 epochs with a ﬁxed learning rate of 1 e − 3 , Adam optimizer (Kingma & Ba, 2014) and a batch size of 50 using the Algorithm 2 from Jabri et al.
Task Policy For task policy π ψ ( s ) , we used a fully connected network with 3 hidden layers of size 256 and ReLU activation.
It then projects the output of the hidden layers to the logits (corresponding to the components of discrete latent space) through linear layers.
CQL Hyperparameters We used the standard hyperparameters for CQL( H ) with discrete action space, as mentioned in Kumar et al.