Title: Parameterizing Non-Parametric Meta-Reinforcement Learning Tasks via Subtask Decomposition
Venue: NeurIPS 2023
Accepted: yes
Keywords: Deep reinforcement learning

Meta-reinforcement learning

Subtask decomposition
Abstract: Meta-reinforcement learning (meta-RL) techniques have demonstrated remarkable success in generalizing deep reinforcement learning across a range of tasks. Nevertheless, these methods often struggle to generalize beyond tasks with parametric variations. To overcome this challenge, we propose Subtask Decomposition and Virtual Training (SDVT), a novel meta-RL approach that decomposes each non-parametric task into a collection of elementary subtasks and parameterizes the task based on its decomposition. We employ a  Gaussian mixture VAE to meta-learn the decomposition process, enabling the agent to reuse policies acquired from common subtasks. Additionally, we propose a virtual training procedure, specifically designed for non-parametric task variability, which generates hypothetical subtask compositions, thereby enhancing generalization to previously unseen subtask compositions. Our method significantly improves performance on the Meta-World ML-10 and ML-45 benchmarks, surpassing current state-of-the-art techniques.

Parameterizing Non-Parametric Meta-Reinforcement Learning Tasks via Subtask Decomposition

Abstract

Meta-reinforcement learning (meta-RL) techniques have demonstrated remarkable success in generalizing deep reinforcement learning across a range of tasks. Nevertheless, these methods often struggle to generalize beyond tasks with parametric variations. To overcome this challenge, we propose Subtask Decomposition and Virtual Training (SDVT), a novel meta-RL approach that decomposes each nonparametric task into a collection of elementary subtasks and parameterizes the task based on its decomposition. We employ a Gaussian mixture VAE to meta-learn the decomposition process, enabling the agent to reuse policies acquired from common subtasks. Additionally, we propose a virtual training procedure, specifically designed for non-parametric task variability, which generates hypothetical subtask compositions, thereby enhancing generalization to previously unseen subtask compositions. Our method significantly improves performance on the Meta-World ML-10 and ML-45 benchmarks, surpassing current state-of-the-art techniques.

1 Introduction

Meta-reinforcement learning (meta-RL) constitutes a dynamic field within deep reinforcement learning, focusing on training agents to quickly adapt to novel tasks by learning from a variety of training tasks [2]. By interacting with these tasks, meta-RL creates an inductive bias regarding the task dynamics and subsequently develops a policy based on this knowledge. Despite its significant contribution to the generalization capability of traditional deep RL, meta-RL is susceptible to test-time distribution shifts, which restricts its applicability to familiar in-distribution test tasks [17, 36, 39, 41].

To tackle this limitation, recent out-of-distribution (OOD) meta-RL approaches have emphasized distinct training and test task distributions, thereby achieving enhanced performance on unseen OOD test tasks with interpolated or slightly extrapolated training dynamics [12, 36, 41, 35, 1]. Although the parameters of training and test tasks are drawn from disjoint distributions, these tasks remain qualitatively similar, as they can be expressed in a shared parametric form representing the task dynamics (e.g., the same 'Pick-place' task with OOD goal positions in Figure 1b).

In this study, we explore a more general meta-RL framework that addresses non-parametric task variability [69, 71], a topic that has received limited attention in prior research. In this context as of Figure 1c, variations among tasks cannot be expressed through simple parametric variations, such as the parameterization of a goal position. Generalization is particularly challenging in this setting, as conventional meta-RL methods often model the inductive bias as a parametric embedding applicable to various tasks [9, 46, 74]. Within a non-parametric framework, it may not be feasible to employ a unified and generalizable parameterization of training tasks using standard meta-RL techniques.

(a) 'Pick-place' in standard indistribution meta-RL setup.

(b) 'Pick-place' in out-ofdistribution meta-RL setup.

(c) Non-parametric task variation between 'Pick-place' and 'Sweep-into.'

Moreover, even if an agent successfully models the inductive bias parametrically, it is improbable that the same parameterization will be reusable for qualitatively distinct test tasks.

In addressing the challenges of non-parametric task variability in meta-RL, our primary strategy involves decomposing each non-parametric task into a set of shared elementary subtasks . We then parameterize each task based on the types of subtasks that constitute it. Despite the non-parametric task variability, tasks may share elementary subtasks. For instance in Figure 1c, a 'Pick-place' task can be decomposed into subtasks: 'grip object' and 'place object,' while a 'Sweep-into' task can be decomposed into subtasks: 'grip object' and 'push object.' By employing the shared subtask parameterization, the policy can capitalize on the captured commonalities between non-parametric tasks to enhance training efficiency and generalization capabilities.

However, our approach of task parameterization based on a set of subtasks faces two primary challenges: the lack of prior information about (1) the set of elementary subtasks and (2) the decomposition of each task. To address these issues, we employ meta-learning for the subtask decomposition (SD) process using a Gaussian mixture variational autoencoder (GMVAE) [29, 8, 59]. Our GMVAE encodes the trajectory up to the current timestep into latent categorical and Gaussian contexts, which are trained to reconstruct the task's reward and transition dynamics [74]. We discover that the meta-learned latent categorical context effectively represents the subtask compositions of tasks under non-parametric variations. Consequently, the policy, using the learned subtask composition, can readily generalize to new tasks comprising previously encountered subtasks. To further enhance generalization to unseen compositions of familiar subtasks, we propose a virtual training (VT) process [35, 1] specifically designed for non-parametric task variability. We train the policy on imaginary tasks generated by the learned dynamics decoder, conditioned on hypothetical subtask compositions.

We evaluate our method on the Meta-World ML-10 and ML-45 benchmarks [71], widely used meta-RL benchmarks comprising diverse non-parametric robotic manipulation tasks. We empirically demonstrate that our method successfully meta-learns the shareable subtask decomposition. With the help of the subtask decomposition and virtual training, our method, without any offline demonstration or test-time gradient updates, achieves test success rates of 33.4% on ML-10 and 31.2% on ML-45, which improves the previous state-of-the-art by approximately 1.7 times and 1.3 times, respectively.

2 Background

2.1 Meta-Reinforcement Learning

A Markov decision process (MDP), M = ( S , A , R, T, T 0 , γ, H ) , is defined by a tuple comprising a set of states S , a set of actions A , a reward function R ( r t +1 | s t , a t , s t +1 ) , a transition function T ( s t +1 | s t , a t ) , an initial state distribution T 0 ( s 0 ) , a discount factor γ , and a horizon H .

The goal of meta-RL is to learn to adapt to a distribution of MDPs with varying reward and transition dynamics. At the start of each meta-training iteration, an MDP is sampled from the distribution p ( M train ) over a set of MDPs M train. Each MDP M k = ( S , A , R k , T k , T 0 ,k , γ, H ) is defined with a unique reward function R k , transition function T k , and initial state distribution T 0 ,k . Unlike in a multi-task setup, the agent in the meta-RL setup does not have access to the task index k that determines the MDP dynamics. The training objective is to optimize the policy π ψ with

parameters ψ to maximize the expected return across all MDPs: max ψ E M k ∼ p ( M train ) [ J pol ( ψ )] , where J pol ( ψ ) = E T 0 ,k ,T k ,π ψ [ ∑ H -1 t =0 γ t R k ( r t +1 | s t , a t , s t +1 ) ] . During meta-testing, standard indistribution meta-RL methods are evaluated on tasks sampled from the same distribution p ( M test ) as the training tasks, i.e., M train = M test = M (Figure 1a). In contrast, OOD meta-RL methods assume strictly disjoint training and test task sets, i.e., M = M train ∪ M test and M train ∩ M test = ∅ (Figure 1b).

2.2 Bayes-adaptive Meta-Reinforcement Learning

Since the true task index k is not provided to the agent in the meta-RL problem setup, it is important to balance exploration and exploitation while learning about the initially unknown MDP. A Bayesadaptive agent [38, 10, 16] achieves this balance by updating its belief b t ( R,T ) about the MDP based on its experience τ : t = { s 0 , a 0 , r 1 , s 1 , a 1 , r 2 , . . . , s t -1 , a t -1 , r t , s t } . The agent's belief over the MDP dynamics at time t can be represented as a posterior given the trajectory, i.e., b t ( R,T ) = p ( R,T | τ : t ) . By augmenting the state with the belief to form a hyper-state space S + = S × B , where B is the set of belief, a Bayes-adaptive MDP (BAMDP) can be constructed. The objective of a BAMDP is to maximize the expected return within a meta-episode while learning, where a meta-episode consists of n roll rollout episodes (i.e., H + = n roll × H steps) of the same MDP:

where T + ( s + t +1 | s + t , a t , r t ) = E b t [ T ( s t +1 | , s t , a t )] I ( b t +1 = p ( R,T | τ : t +1 )) is the transition dynamics and R + ( r t +1 | s + t , a t , s + t +1 ) = E b t +1 [ R ( r t +1 | s t , a t , s t +1 )] is the reward dynamics of the BAMDP. The posterior belief update in the indicator function I ( · ) is intractable for all but simple environments.

VariBAD [74] solves the inference and posterior update of the belief by combining meta-RL and approximate variational inference [28]. At each timestep t , a recurrent encoder q ϕ h encodes the experience τ : t into a hidden embedding h t = q ϕ h ( τ : t ) . The approximate posterior belief over the dynamics can be represented as the parameters of a multivariate Gaussian distribution: b t = ( µ ϕ z ( h t ) , σ ϕ z ( h t )) , where µ ϕ z ( · ) and σ ϕ z ( · ) are neural networks. The latent context z t ∼ N ( µ ϕ z ( h t ) , σ 2 ϕ z ( h t )) is used to estimate the MDP dynamics: p θ R ( r j +1 | s j , a j , s j +1 , z t ) and p θ T ( s j +1 | s j , a j , z t ) for j = 0 , . . . , H + -1 , including the past and future . Then the problem of computing the posterior over the dynamics p ( R,T | τ : t ) reduces to inferring the posterior q ϕ ( z t | τ : t ) , where ϕ = { ϕ h , ϕ z } . A separate policy network π ψ ( a t | s t , b t ) is trained to optimize the BAMDP objective in Eq. (1).

2.3 Non-parametric Task Variability

The term 'non-parametric' in the context of task variability is introduced in the Meta-World paper [69]. It is used to distinguish the task variability in Meta-World manipulation tasks from more simplistic parametric variations exhibited in standard MuJoCo tasks, such as variations in target goal positions, directions, and velocities. However, the term 'non-parametric' might lead to misunderstanding. Because it could suggest that there are no parameters available for task parameterization, even though we have effectively parameterized them in our approach. Our major breakthrough is in enabling this parameterization in terms of subtask compositions, which was difficult using earlier methods that relied on simple latent parametrization. Hence, while we retain the 'non-parametric' terminology, we guide the readers to view the scope of our work through the lens of modularity or composition-based generalization [27].

3 Method

In this section, we introduce our novel meta-RL method, named Subtask Decomposition and Virtual Training (SDVT), to handle non-parametric task variability. Our approach is based on decomposing each task into a set of elementary subtasks and parameterizing each task based on the composition of these subtasks. To achieve this, we use a Gaussian mixture variational autoencoder (GMVAE) to meta-learn the subtask decomposition process. In addition, we introduce a virtual training process that improves generalization to tasks with unseen compositions of subtasks.

Decoder

Policy

~

×

~

×

Dropout

Reconstruction

Encoder

×

Stop gradient

Optional dispersion for SDVT

(not used for SD)

~

(a) SDVT architecture.

0

1

y

(

ht

)

1

0

z

(

ht

,

yt

)

0

100

200

300

400

500

t

5.0

2.5

log

z

(

ht

,

yt

)

1000

2000

3000

4000

5000

t

Pick-place,

K

=5, dim(

z

) = 5

Dim. 1

2

3

4

5

(b) Task inference on 'Pick-place.'

3.1 Subtask Decomposition (SD) with a Gaussian Mixture Variational Autoencoder

Our goal is to meta-learn a set of elementary subtasks and to meta-learn the decomposition of each task into a composition of these subtasks. The core of our method focuses on meta-learning the approximate subtask composition y t ∈ ∆ K sampled from a K -class categorical distribution, where ∆ K denotes the K -dimensional probability simplex. For example with K = 3 , we want the subtask composition y t to be learned somewhat like (0 . 5 , 0 . 5 , 0 . 0) for 'Pick-place' and (0 . 0 , 0 . 5 , 0 . 5) for 'Sweep-into,' where each dimension of y t represents the weight corresponding to subtasks in the order of 'place object,' 'grip object,' and 'push object.' To capture such subtask information, we use a Gaussian mixture variational autoencoder (GMVAE) to represent the latent space of non-parametric tasks as a Gaussian mixture distribution. Each task is represented as a K -dimensional mixture proportion y t , resulting in each class representing a unique subtask shared across different tasks. Thus, the distribution of K subtasks is modeled using a categorical distribution with K classes, where each class is associated with a Gaussian distribution. The learned subtask composition y t represents the agent's belief at time t about the current task's decomposition into subtasks. It is crucial to note that we don't necessarily want y t to be a one-hot embedding representing the subtask that the agent is solving at time t . We want y t to represent a mixing proportion [56] over all possible subtasks that the agent believes to be relevant to the current task, including past and future as in Figures 2b and 4. This distinction is vital to the model's effectiveness in solving a range of tasks, as it allows for flexibility in subtask identification and generalization across different tasks.

Architecture Our full model in Figure 2a, which is based on the VAE of VariBAD [74], consists of three main components: the encoder, decoder, and policy networks parameterized by ϕ = { ϕ h , ϕ y , ϕ z } , θ = { θ z , θ R , θ T } , and ψ , respectively.

(1) The encoder is defined as q ϕ ( y t , z t | h t ) = q ϕ y ( y t | h t ) q ϕ z ( z t | h t , y t ) . A recurrent network encodes the past trajectory τ : t into a hidden embedding h t = q ϕ h ( τ : t ) . First, the categorical encoder q ϕ y ( y t | h t ) : Cat ( ω ϕ y ( h t )) samples y t , where ω ϕ y ( h t ) ∈ ∆ K . We use the Gumbel-Softmax trick [23] with a high temperature ( τ = 1 ) when sampling y t to form a soft label. Then the multivariate Gaussian encoder q ϕ z ( z t | h t , y t ) : N ( µ ϕ z ( h t , y t ) , σ 2 ϕ z ( h t , y t )) samples a continuous latent context z t , which contains the parametric information of the subtasks, in addition to the categorical subtask information of y t .

(2) The decoder is defined as p θ ( τ : H + , y t , z t ) = p ( y t ) p θ z ( z t | y t ) p θ R ,θ T ( τ : H + | z t ) . It reconstructs the reward and transition dynamics for all transitions in the meta-episode τ : H + , using the latent context z t as in Eq. (4). We assume a uniform prior of subtask composition p ( y t ) : Uniform (1 /K ) and a Gaussian regularization p θ z ( z t | y t ) : N ( µ θ z ( y t ) , σ 2 θ z ( y t )) . This encoder-decoder architecture

allows both the approximate posterior q ϕ ( y t , z t | h t ) and the prior p ( z t ) to follow Gaussian mixture distributions.

(3) The policy network, π ψ ( a t | s t , b t ) , is trained separately conditioned on the belief b t = ( µ ϕ z ( h t , y t ) , σ ϕ z ( h t , y t )) that are parameters of the Gaussian context. In practice, we also provide the parameters of the categorical encoder ω ϕ y ( h t ) to the policy, which we find to improve the performance. The parameters of the distributions ( ω ϕ y , µ ϕ z , σ ϕ z , µ θ z , and σ θ z ) are modeled as outputs of multilayer perceptrons (MLPs) as in Appendix C.4.

Objective We optimize the GMVAE to maximize the evidence lower bound (ELBO) for all time steps t = 0 , . . . , H + over the trajectory distribution d ( M k , τ : H + ) induced by the policy in MDP M k :

In addition to the reconstruction objectives J R-rec and J T-rec, we have additional regularization J reg and categorical J cat objectives:

The regularization objective J reg minimizes the KL divergence between the learned posterior Gaussian distribution q ϕ z ( z t | h t , y t ) and learned Gaussian priors p θ z ( z t | y t ) . Unlike the standard VAE that assumes a standard normal prior, we learn K distinct Gaussian priors conditioned on y t . The categorical objective J cat maximizes the conditional entropy of y t given h t and prevents collapse. Refer to Appendix B for the derivation of the ELBO objective. The reconstruction objectives in Eq. (4) are computed for all timesteps from the first to the terminal step of the meta-episode. Therefore, the subtask composition y t at time t is not necessarily a one-hot label of the belief about the current subtask at time t , but a mixture label of the belief on all subtasks that compose the current task in the past and future within the meta-episode. Under the BAMDP setup, the agent learns to reduce the uncertainty of the decomposition as quickly as possible, supporting the policy with the converged parameterization. Combining the policy objective in Eq. (1) and the sum of the ELBO objectives in Eq. (2) for all timesteps in a meta-episode, the overall objective over training tasks is to maximize:

Occupancy regularization The dimension of y t or the number of underlying subtasks K is a crucial hyperparameter that should be determined based on the number of training tasks N train and their similarities. However, in many cases, prior information about the optimal value of K , denoted as K ∗ , may not be available. One way to expand the scope of our method for unknown K ∗ is to meta-learn the number of effective subtasks as well. First, we assume K ∗ < N train, otherwise each task will be classified into a separate subtask with one-hot label, preventing learning shareable subtasks. We start with a sufficiently large K = N train and regularize the ELBO objective to progressively reduce the number of effective subtasks (non-zero components) occupied in y t with the following occupancy regularization that penalizes the usage of larger indices in the subtask composition:

We calculate the dot product of exponential weights and the subtask composition y t to penalize the occupancy of higher dimensions of y t . We scale the dot product by log K to match the scale to the upper bound of J cat. We add J occ multiplied by a coefficient α o to the GMVAE objective in Eq. (3). Consequently, the agent prioritizes using lower indices in the decomposition to represent frequent subtasks and sparingly uses higher indices for rare subtasks as in Figure 4b and in Appendix E.1.

Decoder dropout As the GMVAE is optimized using the trajectories induced by the policy, the decoder can easily overfit the frequent states and actions of training tasks [35]. This can lead to low predicted rewards for unexperienced states and actions, regardless of the latent context z t . When

~

Decoder

~

×

Policy

×

×

Encoder

Dropout

×

Stop gradient

Imaginary subtask composition

Maximize

(a) Imaginary reward generation.

(b)

˜

y

= [1

,

0

,

0

,

0

,

0]

.

(c)

˜

y

= [0

,

1

,

0

,

0

,

0]

.

(d)

˜

y

= [

1

2

,

1

2

,

0

,

0

,

0]

.

(e)

˜

y

= [0

,

0

,

1

3

,

1

3

,

1

3

]

.

rt+1

0.1

f0.10

0.0

0.00

0.8

0.8

0.0

x-axis 0.2

0.1

0.0

0.2

0.2

0.0

0.2

trt+1

at

0.0

0.0

0.8

0.8

0.0

0.05

x-axis

0.2

0.4

0.0

0.2

x-axis 0.05

0.00

0.0

0.2

such overfitting happens, the latent context loses its task-informative value, leading to the potential underperformance of policy based on this context. Following the approach of LDM [35], we address this by applying dropout (DO) to the state and action embeddings of the decoder while leaving latent context z t untouched: p θ R ( r j +1 | DO [ s j , a j , s j +1 ] , z t ) and p θ T ( s j +1 | DO [ s j , a j ] , z t ) . This dropout application is crucial for the virtual training discussed in the following section.

3.2 Virtual Training (VT) on Generated Tasks with Imaginary Subtask Compositions

The overall objective in Eq. (6) is optimized over the trajectories of training tasks. To enhance generalization to test tasks with unseen subtask compositions, we generate virtual training tasks using the imaginary dynamics produced by the GMVAE decoder. This process resembles those in [35, 1], but their generated tasks are limited to parametric variations, e.g., generating tasks with unseen goal positions. Such methods can prepare for test tasks with unseen parametric changes but struggle to prepare for qualitatively new tasks with unseen compositions of subtasks. Our GMVAE model enables us to extend the process to the non-parametric setup by conditioning the decoder on an imaginary subtask composition ˜ y . At the beginning of each meta-episode, we randomly determine with probability p v whether to convert it into a virtual meta-episode. By training the policy on imaginary tasks, it can better prepare for test tasks with unseen subtask compositions in advance. We use the tilde accent to denote imaginary components, and the hat accent to denote estimates.

Latent context dispersion Inspired by the work of Ajay et al. [1], we utilize the dispersion structure for our GMVAE to support extrapolated generalization of the latent context z t . Instead of directly using z t as the decoder input, we insert an additional MLP p θ ˆ h before the decoder to expand the dimension of the context (the dotted box in Figure 2a). The MLP output ˆ h t = p θ ˆ h ( z t ) is trained to reconstruct the embedding h t by appending α d J dis = -α d ∥ h t -ˆ h t ∥ 2 2 to the total GMVAE objective in Eq. (3). We then use the dispersed context ˆ h t instead of z t to reconstruct the reward and transition. This trick is effective in generating imaginary tasks featuring extrapolated dynamics, albeit at the cost of increased training complexity.

Imaginary reward generation Figure 3a presents the imaginary reward generation process. The goal is to create imaginary tasks based on the distribution of training subtasks but with unseen compositions of subtasks. At the beginning of a virtual meta-episode, we randomly sample an imaginary subtask composition ˜ y ∼ Dirichlet (¯ y ) fixed for that virtual meta-episode, where the concentration parameter ¯ y ∈ ∆ K is the empirical running mean of all y t over training. By replacing the real y t with the imaginary subtask composition ˜ y , we sample an imaginary context

nary reward ˜ r t +1 ∼ p θ R ( r t +1 | DO [ s t , a t , s t +1 ] , ˜ ˆ h t ) accordingly using our GMVAE. We replace the reward for the next timestep, r t +1 , with ˜ r t +1 , while the states remain to be from the real training task. The imaginary reward is used for the encoder input at the next time step and for the policy, where the policy is trained to maximize the sum of generated rewards, ∑ H + -1 t =0 γ t ˜ r t +1 . However, the GMVAE is not trained to reconstruct the imaginary dynamics.

3.3 Summary of the Combined Methods: SDVT-LW and SDVT

By combining the subtask decomposition (SD) and virtual training (VT) processes, we propose two methods: SDVT-LW and SDVT. Foremost, our primary contribution is the proposal of SDVT-LW, which is the lightweight (-LW) version of our method that assumes the prior knowledge of the optimal number of subtasks K , therefore not employing the occupancy regularization.

Furthermore, we propose SDVT with the occupancy regularization strategy. This generalizes SDVTLW to adapt to more difficult conditions where there is a lack of prior knowledge of the optimal number of subtasks K ∗ . We initialize the number of subtasks equal to the number of training tasks and employ occupancy regularization to downscale higher dimensions, navigating to discover the most efficient number of subtasks even without prior knowledge.

For the purposes of virtual training, we adopt two methodologies that have found application in previous studies: dropout [35] and dispersion via structured V AE [1]. Their use in our work remains unchanged from their original applications. The efficacy of these components, within the context of virtual training, is demonstrated via ablations in Appendix E. We summarize the entire meta-training process with a pseudocode in Appendix A.

4 Related Work

Classical meta-RL Classical meta-RL methods assume that both training and test tasks are sampled from the same distribution. These methods are divided into two main categories: gradient-based methods and context-based methods. Gradient-based methods [13, 54, 48, 73, 53] learn a common initialization of a lightweight model for all tasks, allowing the agent to achieve high performance on unseen target tasks with a few steps of gradient updates. However, these methods lack online adaptation capability within a meta-episode because they require many pre-update rollouts before adaptation. Context-based methods [9, 18, 46, 32, 34, 43, 74, 40] use a recurrent or memoryaugmented network to encode the collected experience into a latent context. In general, the context is trained to optimize auxiliary objectives such as reward dynamics, transition dynamics, and value function. These methods can adapt to target tasks through online, in-context learning without requiring gradient updates. However, they are vulnerable to test-time distribution shifts since the encoded context and the policy given the context are hardly generalized to out-of-distribution tasks.

Out-of-distribution meta-RL A group of recent studies focuses on training a generalizable agent that is robust to test-time distribution shifts. A group of works generates imaginary observations using image augmentation techniques [18, 21, 31, 44, 33, 61, 66]. Most of these methods depend on predefined heuristic augmentations, without utilizing the training task dynamics. On the other hand, some works explicitly address varying environment dynamics. For example, MIER [7] reuses the trajectories collected during training by relabeling according to the test dynamics. Our work is related to AdMRL [36], LDM [35], and DiAMetR [1], which generate imaginary tasks with unseen dynamics using learned models. However, these works focus on generating parametric variants of training tasks, while we focus on generalizing across non-parametric task variants.

Skill-based meta-RL Our task parameterization based on the subtask decomposition is related to the recently spotlighted skill-based meta-RL methods [11, 51, 50, 55], which aim to achieve fast generalization on unseen tasks by decomposing and extracting reusable skills from training tasks' trajectories. These works often require refined offline demonstrations to learn skills using behavioral cloning objectives, where the skills distinguish a sequence of actions given a sequence of states. For example, SimPL [42] extracts skills from offline demonstrations before meta-training, and during a meta-test, it only adapts the high-level skill policy, with the low-level policy frozen. HeLMS [47] learns a 3-level skill hierarchy by decomposing offline demonstrations of a single task. Online learning of the skills is often unstable because the set of skills should develop along with the online improvement of the policy. The subtask decomposition of our method is conceptually different from

the skill decomposition [68]. Even when the policy is not stationary during online updates, the underlying reward and transition dynamics, which our model has to estimate, do not change.

5 Experiments

5.1 Experimental Setup

Meta-World benchmark The Meta-World V2 benchmark [71] stands as the most prominent, if not the only, established benchmark for assessing meta-RL algorithms featuring non-parametric task variability. This benchmark comprises 50 qualitatively distinct robotic manipulation tasks, with each task containing 50 parametric variants that incorporate randomized goals and initial object positions. Specifically, the Meta-World Meta-Learning 10 (ML-10) benchmark, consists of N train = 10 training tasks and N test = 5 held-out test tasks. We denote each task by an index, where the training tasks are numbered from 1 to 10 and the test tasks from 11 to 15. Likewise, the ML-45 benchmark consists of N train = 45 training tasks and N test = 5 test tasks. Refer to the tables in Appendix D.1 for the set and indices of tasks. The agent must maximize its return from experience while exploring to identify the initially unknown task dynamics within a meta-episode of H + = 5000 steps that consists of n roll = 10 rollout episodes of horizon H = 500 steps each.

SDVT variants and baselines setup We evaluate our methods SDVT and SD (only subtask decomposition without virtual training and dispersion). Without the prior knowledge of K ∗ , we set K = N train and apply the occupancy regularization ( α o = 1 . 0) with α c = 1 . 0 . We also evaluate a lightweight (-LW) version of ours with smaller K = 5 with α c = 0 . 5 and without the occupancy regularization ( α o = 0 . 0 ). To ensure a fair comparison and to exclude the gains from orthogonal contributions, we compare SDVT with state-of-the-art meta-RL methods that do not require any refined offline demonstrations, ensembles, or extensive test-time training: RL 2 [9], MAML [13], PEARL [46], and VariBAD [74]. We also compare with a parametric OOD meta-RL method, LDM [35], to evaluate the efficacy of our virtual training over subtasks. All methods do not perform gradient updates during the test except for MAML. Appendix C presents more implementation details. Briefly, SDVT without a Gaussian mixture reduces to LDM, LDM without virtual training reduces to VariBAD, and VariBAD without a VAE decoder reduces to RL 2 . Our implementation is available at https://github.com/suyoung-lee/SDVT .

Evaluation metric We follow the standard success criterion of Meta-World as follows. A timestep is considered successful if the distance between the task-relevant object and the goal is less than a predefined threshold. A rollout episode is considered successful if the agent ever succeeds at any timestep during the rollout episode. The success of a meta-episode is defined as the success of the last (10th) rollout episode. In Table 1, we report the success rates of 750 (15 tasks × 50 parametric variants) meta-episodes at 250M steps for ML-10 and 2500 (50 tasks × 50 parametric variants) meta-episodes at 400M steps for ML-45. 2 Likewise, we report the returns of the last rollout episodes.

5.2 Results

Table 1 illustrates that no method attains a 100% success rate even on training tasks, emphasizing the challenge posed by non-parametric task variability. The training success rates on ML-45 are consistently lower than those on ML-10, reflecting the inherent difficulty in adapting to a broader range of tasks, such as conflicting gradients [70]. Notably, SD surpasses all other baselines in training success rate. In particular, outperforming VariBAD underscores the limitations of employing a single Gaussian distribution to model the latent task space in cases of non-parametric variability.

On the test tasks, SDVT and SDVT-LW substantially outperform all baselines, even outperforming LDM, which surpasses VariBAD with parametric virtual training. Our gain is attributed to our virtual training process, which is specifically designed for test tasks involving non-parametric variability. Notably, SDVT and SD outperform their LW counterparts in training success rates, primarily due to its fine-grained subtask decomposition which provides a more precise representation of each task. For example, a 'grip object' subtask may be split and represented as a combination of two distinct subtasks 'move gripper' and 'tighten gripper' with a larger K . In contrast, SDVT-LW scores higher

test success rates than SDVT, presumably due to its coarser decomposition that allows imaginary compositions to encompass a broader range of unseen test tasks.

Please refer to the tables in Appendix D.1 for individual task results. Our methods achieve the highest success rates across all test tasks on the ML-10 benchmark. However, all methods encounter challenges in solving 'Shelf-place,' which includes an unseen shelf not present in the observation. As such, this task cannot be decomposed into previously seen subtasks but rather into unseen ones, making it difficult to prepare through virtual training. These tasks, composed of unseen subtasks, pose a considerable challenge for zero-shot adaptation as with SDVT. Addressing these tasks may require a substantial increase in rollouts and updates during tests, an area of potential future work. Detailed ablation results are reported in Appendix E. For additional results such as rendered videos, refer to our webpage https://sites.google.com/view/sdvt-neurips .

5.3 Analysis

In this subsection, we explore whether the performance enhancements attributed to our methods are indeed the result of effectively decomposed subtasks and a diverse range of imaginary tasks.

Learned Subtask Compositions From Figure 2b, we observe a rapid context convergence within the first rollout episode for a sufficiently meta-learned task. To validate our motivation, we visualize the converged subtask compositions in Figure 4. We find that such converged subtask compositions are shared by qualitatively similar tasks. For example, in Figure 4a, tasks (3) 'Pick-place,' (7) 'Peg-insert-side,' and (10) 'Basketball,' which require placing an object at a goal position, share subtask indices 1 and 3. Additionally, simple tasks that require moving the gripper in a straight line: (1) 'Reach,' (5) 'Drawer-close,' and (8) 'Window-open' are classified into the same subtask 2. These observations suggest that our model can effectively meta-learn the decomposition process of tasks into shareable subtasks. Furthermore, the test tasks may not necessarily have the same subtask compositions as the training tasks. For instance, (3) 'Pick-place' and (14) 'Sweep-into' share subtask 3, but not subtasks 1 or 5, revealing the potential for virtual training to be effective.

Occupancy Regularization Figure 4 indicates that it is important to set the subtask class dimension K and the categorical coefficient α c appropriately according to the number of training tasks and the correlations among them. Otherwise, the decomposition may suffer from a collapse or degeneration. With the occupancy regularization, we can avoid the extensive search for the optimal subtask dimension K ∗ and corresponding α c . In Figure 4b, we find that our occupancy regularization successfully limits the use of unnecessary subtasks with higher indices as intended.

Generated imaginary tasks To demonstrate the dynamics of the imaginary tasks, we show the trajectory of the gripper and object over a meta-episode (5000 steps) on generated imaginary tasks for different imaginary subtask compositions ˜ y in Figure 3. When we set ˜ y as a one-hot vector, the policy

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

5

4

3

2

1

yH

+ [ ]

0.000.290.000.000.001.000.000.001.000.001.000.000.000.270.00

0.000.220.001.000.000.000.040.000.000.000.001.000.000.000.00

0.000.480.190.000.000.000.060.000.000.890.000.000.000.730.45

1.000.000.000.001.000.000.001.000.000.000.000.000.000.000.00

0.000.010.810.000.000.000.900.000.000.110.000.001.000.000.55

3. Pick-place

7. Peg-insert-side

10. Basketball

5

4

3

2

1

0.000.000.00

0.000.040.00

0.190.060.89

0.000.000.00

0.810.900.11

1. Reach

5. Drawer-close

8. Window-open

5

4

3

2

1

0.000.000.00

0.000.000.00

0.000.000.00

1.001.001.00

0.000.000.00

(a) SDVT-LW

(

K

= 5

, α

c

= 0

.

5

, α

o

= 0

.

0)

.

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

10

9

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00 0.00 0.00 0.10 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.05 0.00 0.00 0.04

0.00 0.00 0.00 0.30 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.07 0.00 0.00 0.17

0.00 0.00 0.05 0.00 0.00 0.00 0.15 0.26 0.29 0.00 0.00 0.01 0.16 0.01 0.00

0.00 0.00 0.00 0.00 1.00 0.00 0.00 0.11 0.00 0.00 0.00 0.00 0.00 0.00 0.00

1.00 0.09 0.15 0.00 0.00 0.00 0.02 0.00 0.00 0.13 0.01 0.10 0.04 0.35 0.00

0.00 0.00 0.07 0.00 0.00 0.00 0.19 0.31 0.35 0.00 0.00 0.01 0.20 0.01 0.00

0.00 0.00 0.13 0.30 0.00 0.50 0.19 0.00 0.01 0.00 0.48 0.28 0.17 0.04 0.40

0.00 0.00 0.08 0.00 0.00 0.00 0.20 0.32 0.33 0.00 0.00 0.01 0.21 0.02 0.00

0.00 0.91 0.39 0.00 0.00 0.00 0.05 0.00 0.00 0.87 0.01 0.20 0.05 0.52 0.00

0.00 0.00 0.13 0.30 0.00 0.50 0.19 0.00 0.02 0.00 0.50 0.26 0.17 0.04 0.39

Seed 0 at 250M steps

(b) SDVT

(

K

= 10

, α

c

= 1

.

0

, α

o

= 1

.

0)

.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

1.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 1.00 1.00 1.00 0.00 1.00 1.00 1.00 1.00 0.70 0.00 0.00 0.00

0.00 0.00 1.00 0.00 0.00 0.00 1.00 0.00 0.00 0.00 0.00 0.30 1.00 0.90 1.00

0.00 1.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.10 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

(c) SDVT-LW

(

K

= 5

, α

c

= 0

.

1

, α

o

= 0

.

0)

.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

0.26 0.17 0.18 0.13 0.27 0.26 0.20 0.22 0.16 0.20 0.24 0.19 0.20 0.18 0.16

0.07 0.26 0.24 0.33 0.05 0.07 0.20 0.15 0.27 0.20 0.11 0.20 0.19 0.22 0.27

0.23 0.19 0.19 0.17 0.23 0.23 0.20 0.21 0.18 0.20 0.22 0.20 0.20 0.19 0.18

0.25 0.18 0.19 0.16 0.25 0.25 0.20 0.21 0.18 0.20 0.23 0.20 0.20 0.19 0.18

0.19 0.21 0.21 0.21 0.19 0.20 0.21 0.20 0.21 0.21 0.20 0.21 0.21 0.21 0.21

(d) SDVT-LW

(

K

= 5

, α

c

= 2

.

0

, α

o

= 0

.

0)

.

returns a trajectory that solves an elementary subtask, such as placing and reaching up. We observe that the trajectories vary across different ˜ y , and similar compositions result in similar trajectories.

6 Conclusion

In conclusion, our proposed method has demonstrated a considerable enhancement in meta-RL with non-parametric task variability. This improvement is achieved by meta-learning shareable subtask decompositions and executing virtual training on the imaginary subtask compositions. However, it's essential to acknowledge certain potential limitations beyond the scope of this study, particularly when addressing test tasks involving entirely novel subtasks and in broader setups where the action and state spaces may also vary. Despite these limitations, we believe that expanding the realm of meta-RL to accommodate a wider range of task variability is a critical research topic. Incorporating orthogonal approaches such as using offline demonstrations or test time training techniques into our method could lead to interesting future work addressing the limitations. We are optimistic that our study lays a robust foundation for future research in this field.

References

A. Ajay, A. Gupta, D. Ghosh, S. Levine, and P. Agrawal. Distributionally adaptive meta reinforcement learning. In Neural Information Processing Systems (NeurIPS) , 2022.

J. Beck, R. Vuorio, E. Z. Liu, Z. Xiong, L. Zintgraf, C. Finn, and S. Whiteson. A survey of meta-reinforcement learning. arXiv preprint arXiv: 2301.08028 , 2023.

A. R. Cassandra, L. P. Kaelbling, and M. L. Littman. Acting optimally in partially observable stochastic domains. In National Conference on Artificial Intelligence (NCAI) , 1994.

Y. Chandak, G. Theocharous, S. Shankar, M. White, S. Mahadevan, and P. Thomas. Optimizing for the future in non-stationary mdps. In International Conference on Machine Learning (ICML) , pages 119:1414-1425, 2020.

Y. Chebotar, A. Handa, V. Makoviychuk, M. Macklin, J. Issac, N. D. Ratliff, and D. Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In International Conference on Robotics and Automation, (ICRA) , pages 8973-8979, 2019.

B. Chen, W. Deng, and H. Shen. Virtual class enhanced discriminative embedding learning. In Neural Information Processing Systems (NeurIPS) , pages 31:1942-1952, 2018.

K. Cobbe, O. Klimov, C. Hesse, T. Kim, and J. Schulman. Quantifying generalization in reinforcement learning. In International Conference on Machine Learning (ICML) , pages 97:1282-1289, 2019.

N. Dilokthanakul, P. A. Mediano, M. Garnelo, M. C. Lee, H. Salimbeni, K. Arulkumaran, and M. Shanahan. Deep unsupervised clustering with gaussian mixture variational autoencoders. In International Conference on Learning Representations (ICLR) , 2017.

Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel. RL 2 : Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1611.02779 , 2016.

M. O. Duff. Optimal Learning: Computational Procedures for Bayes-adaptive Markov Decision Processes . University of Massachusetts at Amherst, 2002.

B. Eysenbach, A. Gupta, J. Ibarz, and S. Levine. Diversity is all you need: Learning skills without a reward function. In International Conference on Learning Representations (ICLR) , 2019.

R. Fakoor, P. Chaudhari, S. Soatto, and A. J. Smola. Meta-q-learning. In International Conference on Learning Representations (ICLR) , 2020.

C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning (ICML) , pages 70:1126-1135, 2017.

C. Florensa, D. Held, X. Geng, and P. Abbeel. Automatic goal generation for reinforcement learning agents. In International Conference on Machine Learning (ICML) , pages 80:15151528, 2018.

Garage-Contributors. Garage: A toolkit for reproducible reinforcement learning research. https://github.com/rlworkgroup/garage , 2019.

M. Ghavamzadeh, S. Mannor, J. Pineau, and A. Tamar. Bayesian reinforcement learning: A survey. Found. Trends Mach. Learn. , pages 8(5-6):359-483, 2015.

A. Gupta, B. Eysenbach, C. Finn, and S. Levine. Unsupervised meta-learning for reinforcement learning. arXiv preprint arXiv: 1806.04640 , 2018.

A. Gupta, R. Mendonca, Y. Liu, P. Abbeel, and S. Levine. Meta-reinforcement learning of structured exploration strategies. In Neural Information Processing Systems (NeurIPS) , pages 31:5302-5311, 2018.

D. Hafner, T. P. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, (ICLR) , 2020.

J. Humplik, A. Galashov, L. Hasenclever, P. A. Ortega, Y. W. Teh, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: 1905.06424 , 2019.

M. Igl, K. Ciosek, Y. Li, S. Tschiatschek, C. Zhang, S. Devlin, and K. Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. In Neural Information Processing Systems (NeurIPS) , pages 32:13978-13990, 2019.

A. Jabri, K. Hsu, B. Eysenbach, A. Gupta, S. Levine, and C. Finn. Unsupervised curricula for visual meta-reinforcement learning. In Neural Information Processing Systems (NeurIPS) , volume 32, 2019.

E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations (ICLR) , 2017.

J. Kaddour, S. Sæmundsson, and M. P. Deisenroth. Probabilistic active meta-learning. In Neural Information Processing Systems (NeurIPS) , pages 33:20813-20822, 2020.

L. P. Kaelbling, M. L.Littman, and A. R.Cassandra. Acting optimally in partially observable stochastic domains. In National Conference on Artificial Intelligence , 1994.

L. P. Kaelbling, M. L.Littman, and A. R.Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence , 101:99-134, 1998.

K. Khetarpal, M. Riemer, I. Rish, and D. Precup. Towards continual reinforcement learning: A review and perspectives. In Journal of Artificial Intelligence Research , pages 1401-1476, 2022.

D. P. Kingma and M. Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR) , 2014.

D. P. Kingma, D. J. Rezende, S. Mohamed, and M. Welling. Semi-supervised learning with deep generative models. In Neural Information Processing Systems (NeurIPS) , 2014.

L. Kirsch, S. van Steenkiste, and J. Schmidhuber. Improving generalization in meta reinforcement learning using learned objectives. In International Conference on Learning Representations (ICLR) , 2020.

M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data. In Neural Information Processing Systems (NeurIPS) , pages 33:19884-19895, 2020.

A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. In Neural Information Processing Systems (NeurIPS) , pages 33:741-752, 2020.

K. Lee, K. Lee, J. Shin, and H. Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In International Conference on Learning Representations (ICLR) , 2020.

K. Lee, Y. Seo, S. Lee, H. Lee, and J. Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning (ICML) , pages 119:5757-5766, 2020.

S. Lee and S. Chung. Improving generalization in meta-RL with imaginary tasks from latent dynamics mixture. In Neural Information Processing Systems (NeurIPS) , pages 34:2722227235, 2021.

Z. Lin, G. Thomas, G. Yang, and T. Ma. Model-based adversarial meta-reinforcement learning. In Neural Information Processing Systems (NeurIPS) , pages 33:10161-10173, 2020.

S. Liu, A. J. Davison, and E. Johns. Self-supervised generalisation with meta auxiliary learning. In Neural Information Processing Systems (NeurIPS) , pages 32:1679-1689, 2019.

J. J. Martin. Bayesian Decision Problems and Markov Chains . Wiley, 1967.

B. Mehta, T. Deleu, S. C. Raparthy, C. J. Pal, and L. Paull. Curriculum in gradient-based meta-reinforcement learning. arXiv preprint arXiv:2002.07956 , 2020.

L. C. Melo. Transformers are meta-reinforcement learners. In International Conference on Machine Learning (ICML) , pages 162:15340-15359, 2022.

R. Mendonca, X. Geng, C. Finn, and S. Levine. Meta-reinforcement learning robust to distributional shift via model identification and experience relabeling. arXiv preprint arXiv: 2006.07178 , 2020.

T. Nam, S.-H. Sun, K. Pertsch, S. J. Hwang, and J. J. Lim. Skill-based meta-reinforcement learning. In International Conference on Learning Representations (ICLR) , 2022.

R. Raileanu, M. Goldstein, A. Szlam, and R. Fergus. Fast adaptation to new environments via policy-dynamics value functions. In International Conference on Machine Learning (ICML) , pages 119:7920-7931, 2020.

R. Raileanu, M. Goldstein, D. Yarats, I. Kostrikov, and R. Fergus. Automatic data augmentation for generalization in reinforcement learning. In Neural Information Processing Systems (NeurIPS) , pages 34:5402-5415, 2021.

A. Rajeswaran, K. Lowrey, E. Todorov, and S. M. Kakade. Towards generalization and simplicity in continuous control. In Neural Information Processing Systems (NeurIPS) , pages 30:6550-6561, 2017.

K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International Conference on Machine Learning (ICML) , pages 97:5331-5340, 2019.

D. Rao, F. Sadeghi, L. Hasenclever, M. Wulfmeier, M. Zambelli, G. Vezzani, D. Tirumala, Y. Aytar, J. Merel, N. Heess, and R. Hadsell. Learning transferable motor skills with hierarchical latent mixture policies. In International Conference on Learning Representations (ICLR) , 2022.

J. Rothfuss, D. Lee, I. Clavera, T. Asfour, and P. Abbeel. Promp: Proximal meta-policy search. In International Conference on Learning Representations (ICLR) , 2019.

J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347 , 2017.

D. Shah, P. Xu, Y. Lu, T. Xiao, A. Toshev, S. Levine, and B. Ichter. Value function spaces: Skill-centric state abstractions for long-horizon reasoning. In International Conference on Learning Representations (ICLR) , 2022.

L. X. Shi, J. J. Lim, and Y . Lee. Skill-based model-based reinforcement learning. In Conference on Robot Learning, (CoRL) , 2022.

X. Song, Y. Jiang, S. Tu, Y. Du, and B. Neyshabur. Observational overfitting in reinforcement learning. In International Conference on Learning Representations (ICLR) , 2020.

Y. Song, A. Mavalankar, W. Sun, and S. Gao. Provably efficient model-based policy adaptation. In International Conference on Machine Learning (ICML) , pages 119:9088-9098, 2020.

B. C. Stadie, G. Yang, R. Houthooft, X. Chen, Y. Duan, Y. Wu, P. Abbeel, and I. Sutskever. The importance of sampling in meta-reinforcement learning. In Neural Information Processing Systems (NeurIPS) , pages 31:9280-9290, 2018.

L. Sun, H. Zhang, W. Xu, and M. Tomizuka. Paco: Parameter-compositional multi-task reinforcement learning. In Neural Information Processing Systems (NeurIPS) , 2022.

Y. W. Teh, M. I. Jordan, M. J. Beal, and D. M. Blei. Sharing clusters among related groups: Hierarchical dirichlet processes. In Neural Information Processing Systems (NeurIPS) , 2004.

J. Tobin, R. Fong, A. Ray, J. Schneider, W. Zaremba, and P. Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In International Conference on Intelligent Robots and Systems (IROS) , pages 23-30, 2017.

E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems (IROS) , pages 5026-5033, 2012.

Y. B. Varolgunes, T. Bereau, and J. F. Rudzinski. Interpretable embeddings from molecular simulations using gaussian mixture variational autoencoders. Machine Learning Science and Technology , 1, 2020.

J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick. Learning to reinforcement learn. In Annual Meeting of the Cognitive Science Community (CogSci) , 2016.

K. Wang, B. Kang, J. Shao, and J. Feng. Improving generalization in reinforcement learning with mixture regularization. In Neural Information Processing Systems (NeurIPS) , pages 33:7968-7978, 2020.

R. Wang, J. Lehman, J. Clune, and K. O. Stanley. Paired open-ended trailblazer (POET): endlessly generating increasingly complex and diverse learning environments and their solutions. arXiv preprint arXiv: 1901.01753 , 2019.

R. Wang, J. Lehman, A. Rawal, J. Zhi, Y. Li, J. Clune, and K. Stanley. Enhanced poet: Openended reinforcement learning through unbounded invention of learning challenges and their solutions. In International Conference on Machine Learning (ICML) , pages 119:9940-9951, 2020.

S. Whiteson, B. Tanner, M. E. Taylor, and P. Stone. Protecting against evaluation overfitting in empirical reinforcement learning. In IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL) , pages 120-127, 2011.

J. Yang, B. K. Petersen, H. Zha, and D. Faissol. Single episode policy transfer in reinforcement learning. In International Conference on Learning Representations (ICLR) , 2020.

D. Yarats, I. Kostrikov, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations (ICLR) , 2021.

M. Yin, G. Tucker, M. Zhou, S. Levine, and C. Finn. Meta-learning without memorization. In International Conference on Learning Representations (ICLR) , 2020.

M. Yoo, S. Cho, and H. Woo. Skills regularized task decomposition for multi-task offline reinforcement learning. In Neural Information Processing Systems (NeurIPS) , 2022.

T. Yu, D. Quillen, Z. He, R. Julian, K. Hausman, C. Finn, and S. Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, (CoRL) , 2019.

T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn. Gradient surgery for multi-task learning. In Neural Information Processing Systems (NeurIPS) , pages 33:5824-5836, 2020.

T. Yu, D. Quillen, Z. He, R. Julian, A. Narayan, H. Shively, A. Bellathur, K. Hausman, C. Finn, and S. Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. arXiv preprint arXiv:1910.10897 , 2021.

C. Zhang, O. Vinyals, R. Munos, and S. Bengio. A study on overfitting in deep reinforcement learning. arXiv preprint arXiv: 1804.06893 , 2018.

L. M. Zintgraf, K. Shiarlis, V . Kurin, K. Hofmann, and S. Whiteson. Fast context adaptation via meta-learning. In International Conference on Machine Learning (ICML) , pages 97:7693-7702, 2019.

L. M. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y . Gal, K. Hofmann, and S. Whiteson. Varibad: Avery good method for bayes-adaptive deep RL via meta-learning. In International Conference on Learning Representations (ICLR) , 2020.

A Pseudocode

Algorithm 1 Subtask Decomposition and Virtual Training (SDVT)

Initialize Encoder q ϕ , decoder p θ , policy π ψ , set of training tasks M train, virtual ratio p v , GMVAE buffer B VAE, policy buffer B pol, total number of meta-episodes to train on n meta, number of rollout episodes per meta-episode n roll, running mean of subtask composition ¯ y .

for meta-episode k = 0 , . . . , n meta -1 do Sample a training task M k ∼ p ( M train ) Sample a virtual meta-episode flag V ∼ Bernoulli ( p v ) if V = 1 then Sample an imaginary subtask composition ˜ y ∼ Dirichlet (¯ y ) end if Reset h 0 , y 0 , B pol for timestep t = 0 , . . . , n roll × H -1 do if t mod H = 0 then Reset rollout episode s t ∼ T 0 ,k ( · ) end if Sample an action a t ∼ π ψ ( ·| s t , µ ϕ z ( h t , y t ) , σ ϕ z ( h t , y t ) , ω ϕ y ( h t )) Take an environment step s t +1 ∼ T k ( ·| s t , a t ) r t +1 ← R k ( s t , a t , s t +1 ) if V = 1 then ˜ z t ∼ N ( µ ϕ z ( h t , ˜ y ) , σ 2 ϕ z ( h t , ˜ y )) ˜ ˆ h t = p θ ˆ h (˜ z t ) ˜ r t +1 ∼ p θ R ( ·| DO [ s t , a t , s t +1 ] , ˜ ˆ h t ) Replace r t +1 ← ˜ r t +1 else Add the transition ( s t , a t , s t +1 , r t +1 ) to B VAE · · · VAE NOT TRAINED WITH VIRTUAL DYNAMICS end if Add the transition ( s t , a t , s t +1 , r t +1 ) to B pol Update hidden embedding h t +1 = q ϕ h ( τ : t +1 ) Update subtask composition y t +1 ∼ Cat ( ω ϕ y ( h t +1 ) ) Update the running mean of subtask composition ¯ y ← RunningMeanUpdate (¯ y, y t +1 ) end for Update GMVAE { ϕ, θ } ← { ϕ, θ } + ∇ { ϕ,θ } ∑ H + t =0 ELBO t with samples from B VAE Update policy ψ ← ψ + ∇ ψ J + pol with samples from B pol Anneal virtual ratio p v ← p v +∆ p v end for

B ELBO Derivation

The GMVAE's ELBO objective is derived as follows.

Eq. (8) is equivalent to the ELBO objective in Eq. (3) without weighting coefficients. We assume that the reconstruction τ : H + is conditionally independent of the subtask composition y t given z t .

C Implementation Details

C.1 Reference Implementations

MAML, RL 2 , and PEARL To replicate the results of MAML [13], RL 2 [9], and PEARL [46] as reported in the Meta-World paper [71], we utilize the exact version 3 of the Garage repository [15] without any modifications. For their hyperparameters, please refer to Appendix D.7, D.8, and D.9 of the Meta-World paper. MAML is the only baseline that has the advantage to take gradient updates during the test.

SDVT, SD, LDM, VariBAD Task-inference-based methods (SDVT, SD, LDM [35], and VariBAD [74]) are adapted to the Meta-World benchmark based on the VariBAD's implementation. 4 We begin with VariBAD's hyperparameter configuration for MuJoCo [58] Ant-goal and modify certain hyperparameters to accommodate the Meta-World benchmark (e.g., batch size, network capacity, etc.). In line with VariBAD and LDM, we train SDVT's policy using PPO [49]. The GMVAE is based on an implementation 5 that employs the Gumbel-Softmax reparameterization trick [23] when sampling the categorical subtask variable. The virtual training processes of SDVT and LDM require agent-environment interaction since they use states from the real environment. Therefore, the virtual training steps are also added when counting the total number of training steps. For a comprehensive list of SDVT hyperparameters, shared among SD, LDM, and VariBAD to ensure a fair comparison, please refer to Appendix C.3 and our source code available at https://github. com/suyoung-lee/SDVT .

RL 2

Recurrent encoder

Policy network

VariBAD

VAE

(recurrent encoder &

dynamics decoder)

Policy network

SD

GMVAE

(recurrent encoder &

dynamics decoder)

Policy network

SDVT

GMVAE

(recurrent encoder &

dynamics decoder)

Policy network

Virtual training (subtask)

LDM

VAE

(recurrent encoder &

dynamics decoder)

Policy network

Virtual training (parametric)

+ Dynamics decoder

+ Gaussian mixture

+ Gaussian mixture

+ Virtual training

(parametric)

+ Virtual training

(subtask)

C.2 Computational Complexity

Our experiments were conducted using an Nvidia TITAN Xp. Due to the considerably larger variety of tasks compared to standard meta-RL benchmarks, non-parametric benchmarks require significantly more computational resources and time. In Table 2, we detail the total wall-clock time expended by our methods and baselines during the training and evaluation of the ML-10 benchmark. This time includes the environmental interaction time for 250M training steps, roughly 100M steps dedicated to

evaluations, and the time taken to train the neural networks. Despite incorporating the GMV AE and virtual training, our method's computational demand does not substantially surpass that of VariBAD.

C.3 Hyperparameters

C.4 Network Architecture

The network architectures comprising our method are described in Table 5. Prior to being input into the encoder or decoder, all state, action, and reward inputs pass through embedding networks. The values of K = dim ( y ) and dim ( z ) differ for the ML-10 and ML-45, as indicated in Table 3.

C.5 Aggregation method for the Success Rate

The main results in the Meta-World paper (Table 1 and Figure 6 of Yu et al. [71]) present the average of the maximum success rate for each task, diverging from the raw scores in their Figure 17 and 18. For instance, if an agent achieves a 90% success rate on 'Door-close' in one evaluation during training and scores 10% in all other evaluations at different times, the reported success rate for 'Door-close' is 90%. The mean of these maximum scores across tasks is reported as the aggregated success rate.

This unconventional aggregation method does not accurately represent the meta-RL objective, which aims to train a single agent capable of solving multiple tasks. Meta-World calculates success rates for various tasks at distinct time points during training, even though the agent might specialize in different tasks at various stages. As a result, the agent may excel at specific tasks by chance. Consequently, evaluating the agent more frequently could yield higher maximum scores. We report the conventional final performance, which better reflects the meta-RL objective.

D Detailed Experimental Results

D.1 Performance on Individual Tasks

D.2 Learning Curves

In Figures 6 through 9, the mean training and test success rates are represented by solid and dashed lines respectively, with the shaded region indicating the 95% confidence interval. All results are compiled from 8 random seeds. The data corresponding to the final training steps in these plots (250M steps for ML-10 and 400M steps for ML-45) are used to present the main results in Table 1.

0.0

0.5

1.0

1.5

2.0

2.5

Steps

1e8

0

20

40

60

80

100

Success Rate (%)

SDVT

SDVT-LW

SD

SD-LW

RL2

MAML

PEARL

VariBAD

LDM

0.0

0.5

1.0

1.5

2.0

2.5

Steps

1e8

0

1000

2000

3000

4000

5000

Return

SDVT

SDVT-LW

SD

SD-LW

RL2

MAML

PEARL

VariBAD

LDM

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e8

0

20

40

60

80

100

Success Rate (%)

SDVT

SDVT-LW

SD

SD-LW

RL2

MAML

PEARL

VariBAD

LDM

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Steps

1e8

0

1000

2000

3000

4000

5000

Return

SDVT

SDVT-LW

SD

SD-LW

RL2

MAML

PEARL

VariBAD

LDM

D.3 Subtask Compositions Learned over Training

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.20

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.19

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

0.21

Seed 0 at 0M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

0.00

0.00

0.00

0.45

0.00

0.50

0.00

0.00

0.47

0.10

0.58

0.00

0.00

0.00

0.30

0.00

0.00

0.00

0.38

0.00

0.50

0.00

0.00

0.53

0.20

0.42

0.00

0.00

0.00

0.31

0.26

0.49

0.49

0.00

0.00

0.00

0.49

0.00

0.00

0.36

0.00

0.50

0.49

0.50

0.16

0.74

0.00

0.00

0.16

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.15

0.00

0.51

0.51

0.00

0.00

0.00

0.51

0.00

0.00

0.34

0.00

0.50

0.51

0.50

0.08

Seed 0 at 10M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

0.06

0.14

0.07

0.65

0.00

0.55

0.00

0.00

0.41

0.06

0.61

0.53

0.00

0.01

0.38

0.07

0.21

0.10

0.35

0.00

0.45

0.01

0.00

0.48

0.09

0.39

0.39

0.00

0.02

0.45

0.08

0.36

0.43

0.00

0.00

0.00

0.42

0.00

0.07

0.46

0.00

0.04

0.41

0.39

0.10

0.74

0.02

0.03

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.01

0.00

0.17

0.00

0.04

0.27

0.37

0.00

0.00

0.00

0.56

0.00

0.03

0.39

0.00

0.02

0.59

0.41

0.06

Seed 0 at 20M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

0.00

0.09

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.59

0.16

0.00

0.00

0.00

0.00

0.19

0.01

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.41

0.09

0.00

0.71

1.00

0.00

0.68

0.53

0.00

0.00

0.00

0.27

0.00

0.00

0.98

0.00

0.75

0.21

0.29

0.00

1.00

0.04

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.01

0.00

0.00

0.00

0.00

0.00

0.46

0.00

0.00

0.00

0.73

0.00

0.00

0.02

0.00

0.00

0.79

0.00

0.00

Seed 0 at 50M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Task Index

5

4

3

2

1

yH

+ [ ]

0.00

0.34

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.99

0.01

0.00

0.00

0.01

0.00

0.13

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.01

0.79

0.00

0.00

0.98

0.00

0.53

0.75

0.00

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.21

0.00

0.00

0.02

1.00

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.25

0.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

1.00

1.00

0.00

Seed 0 at 100M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Task Index

5

4

3

2

1

0.00

0.29

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

1.00

0.00

0.00

0.27

0.00

0.00

0.22

0.00

1.00

0.00

0.00

0.04

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.48

0.19

0.00

0.00

0.00

0.06

0.00

0.00

0.89

0.00

0.00

0.00

0.73

0.45

1.00

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.01

0.81

0.00

0.00

0.00

0.90

0.00

0.00

0.11

0.00

0.00

1.00

0.00

0.55

Seed 0 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

10

9

8

7

6

5

4

3

2

1

yH

+ [ ]

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.09

0.10

0.10

0.10

0.10

0.10

0.10

0.09

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

0.10

Seed 0 at 0M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

10

9

8

7

6

5

4

3

2

1

0.00

0.01

0.00

0.02

0.00

0.01

0.00

0.00

0.00

0.00

0.01

0.02

0.00

0.01

0.01

0.00

0.07

0.07

0.10

0.00

0.06

0.07

0.02

0.03

0.02

0.06

0.09

0.07

0.07

0.08

0.03

0.07

0.07

0.06

0.00

0.08

0.07

0.14

0.11

0.12

0.08

0.06

0.07

0.07

0.07

0.95

0.08

0.08

0.07

1.00

0.10

0.08

0.37

0.12

0.13

0.10

0.07

0.08

0.08

0.08

0.01

0.11

0.11

0.08

0.00

0.10

0.11

0.10

0.15

0.17

0.10

0.09

0.11

0.10

0.10

0.00

0.12

0.12

0.13

0.00

0.12

0.12

0.09

0.12

0.12

0.12

0.12

0.12

0.12

0.13

0.00

0.15

0.15

0.16

0.00

0.15

0.15

0.06

0.10

0.09

0.15

0.16

0.15

0.15

0.15

0.01

0.11

0.11

0.10

0.00

0.13

0.11

0.11

0.12

0.12

0.13

0.10

0.11

0.11

0.11

0.00

0.13

0.13

0.12

0.00

0.12

0.13

0.08

0.13

0.13

0.12

0.12

0.13

0.13

0.12

0.00

0.15

0.15

0.18

0.00

0.13

0.15

0.05

0.11

0.10

0.14

0.17

0.15

0.15

0.16

Seed 0 at 10M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

10

9

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00

0.01

0.01

0.01

0.00

0.01

0.01

0.00

0.00

0.00

0.01

0.02

0.01

0.01

0.01

0.01

0.08

0.08

0.10

0.01

0.07

0.08

0.02

0.02

0.01

0.07

0.11

0.08

0.08

0.06

0.15

0.07

0.07

0.05

0.11

0.09

0.07

0.13

0.14

0.16

0.07

0.06

0.07

0.07

0.10

0.42

0.03

0.03

0.02

0.55

0.06

0.04

0.21

0.09

0.08

0.05

0.02

0.03

0.04

0.05

0.08

0.12

0.12

0.08

0.08

0.11

0.12

0.11

0.17

0.22

0.10

0.09

0.11

0.12

0.14

0.10

0.12

0.12

0.11

0.07

0.12

0.12

0.13

0.13

0.13

0.12

0.11

0.12

0.12

0.12

0.03

0.15

0.16

0.20

0.02

0.14

0.15

0.08

0.08

0.06

0.16

0.18

0.16

0.15

0.13

0.12

0.12

0.12

0.11

0.08

0.12

0.12

0.14

0.12

0.11

0.13

0.11

0.12

0.12

0.11

0.07

0.14

0.14

0.12

0.06

0.13

0.14

0.11

0.14

0.14

0.13

0.12

0.14

0.14

0.14

0.02

0.16

0.16

0.19

0.03

0.14

0.16

0.07

0.11

0.10

0.16

0.17

0.16

0.16

0.15

Seed 0 at 20M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

10

9

8

7

6

5

4

3

2

1

0.00

0.00

0.00

0.12

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.02

0.00

0.00

0.01

0.00

0.01

0.00

0.36

0.00

0.11

0.01

0.00

0.00

0.00

0.13

0.09

0.01

0.01

0.13

0.07

0.11

0.11

0.00

0.00

0.00

0.11

0.23

0.12

0.04

0.00

0.05

0.11

0.11

0.02

0.87

0.00

0.00

0.00

1.00

0.00

0.00

0.34

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.18

0.18

0.00

0.00

0.01

0.18

0.04

0.21

0.53

0.00

0.09

0.18

0.18

0.04

0.01

0.18

0.18

0.00

0.00

0.11

0.18

0.12

0.19

0.08

0.09

0.12

0.18

0.18

0.09

0.00

0.09

0.08

0.25

0.00

0.36

0.09

0.01

0.05

0.00

0.38

0.21

0.09

0.09

0.31

0.04

0.16

0.16

0.00

0.00

0.02

0.16

0.22

0.17

0.07

0.01

0.09

0.16

0.16

0.04

0.00

0.19

0.20

0.00

0.00

0.04

0.19

0.03

0.21

0.28

0.02

0.12

0.19

0.19

0.07

0.00

0.08

0.08

0.26

0.00

0.35

0.09

0.01

0.05

0.00

0.37

0.20

0.08

0.09

0.29

Seed 0 at 50M steps

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

10

9

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00

0.00

0.00

0.13

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.02

0.00

0.00

0.11

0.00

0.00

0.00

0.34

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.02

0.00

0.00

0.18

0.10

0.11

0.11

0.00

0.00

0.00

0.11

0.34

0.09

0.00

0.00

0.01

0.11

0.11

0.00

0.83

0.00

0.00

0.00

1.00

0.00

0.00

0.09

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.12

0.12

0.00

0.00

0.00

0.12

0.00

0.14

0.49

0.00

0.26

0.12

0.12

0.08

0.01

0.18

0.19

0.00

0.00

0.00

0.19

0.20

0.19

0.00

0.00

0.02

0.19

0.19

0.00

0.00

0.16

0.16

0.26

0.00

0.48

0.16

0.01

0.16

0.00

0.49

0.19

0.16

0.16

0.26

0.06

0.15

0.15

0.00

0.00

0.00

0.15

0.35

0.13

0.00

0.00

0.01

0.15

0.16

0.00

0.00

0.12

0.11

0.00

0.00

0.00

0.11

0.00

0.13

0.51

0.00

0.27

0.11

0.11

0.11

0.00

0.15

0.15

0.27

0.00

0.52

0.15

0.01

0.16

0.00

0.51

0.19

0.15

0.15

0.26

Seed 0 at 100M steps

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

10

9

8

7

6

5

4

3

2

1

0.00

0.00

0.00

0.10

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.05

0.00

0.00

0.04

0.00

0.00

0.00

0.30

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.07

0.00

0.00

0.17

0.00

0.00

0.05

0.00

0.00

0.00

0.15

0.26

0.29

0.00

0.00

0.01

0.16

0.01

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.11

0.00

0.00

0.00

0.00

0.00

0.00

0.00

1.00

0.09

0.15

0.00

0.00

0.00

0.02

0.00

0.00

0.13

0.01

0.10

0.04

0.35

0.00

0.00

0.00

0.07

0.00

0.00

0.00

0.19

0.31

0.35

0.00

0.00

0.01

0.20

0.01

0.00

0.00

0.00

0.13

0.30

0.00

0.50

0.19

0.00

0.01

0.00

0.48

0.28

0.17

0.04

0.40

0.00

0.00

0.08

0.00

0.00

0.00

0.20

0.32

0.33

0.00

0.00

0.01

0.21

0.02

0.00

0.00

0.91

0.39

0.00

0.00

0.00

0.05

0.00

0.00

0.87

0.01

0.20

0.05

0.52

0.00

0.00

0.00

0.13

0.30

0.00

0.50

0.19

0.00

0.02

0.00

0.50

0.26

0.17

0.04

0.39

Seed 0 at 250M steps

D.4 Subtask Compositions of all Seeds

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

0.00

0.29

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

1.00

0.00

0.00

0.27

0.00

0.00

0.22

0.00

1.00

0.00

0.00

0.04

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.48

0.19

0.00

0.00

0.00

0.06

0.00

0.00

0.89

0.00

0.00

0.00

0.73

0.45

1.00

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.01

0.81

0.00

0.00

0.00

0.90

0.00

0.00

0.11

0.00

0.00

1.00

0.00

0.55

Seed 0 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

0.28

0.42

0.41

0.00

0.00

0.00

0.75

0.00

0.40

0.00

0.00

0.00

0.85

0.49

0.06

0.42

0.01

0.00

1.00

0.00

0.00

0.00

0.00

0.60

0.00

0.00

1.00

0.00

0.10

0.10

0.00

0.00

0.07

0.00

0.00

0.00

0.25

0.00

0.00

1.00

0.00

0.00

0.14

0.13

0.00

0.30

0.57

0.52

0.00

0.00

1.00

0.00

0.00

0.00

0.00

1.00

0.00

0.01

0.28

0.84

0.00

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

Seed 1 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

0.92

1.00

0.13

1.00

0.00

1.00

0.00

0.00

1.00

0.98

1.00

1.00

0.01

0.97

0.94

0.08

0.00

0.87

0.00

0.00

0.00

1.00

0.00

0.00

0.02

0.00

0.00

0.99

0.03

0.06

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.42

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.58

0.00

0.00

0.00

0.00

0.00

0.00

0.00

Seed 2 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

0.47

0.00

0.00

0.00

1.00

0.00

0.00

0.21

0.00

0.00

0.00

0.04

0.00

0.00

0.00

0.00

0.77

1.00

0.00

0.00

0.00

0.98

0.00

0.00

0.00

0.00

0.05

1.00

1.00

1.00

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.93

0.00

0.76

0.73

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.07

1.00

0.24

0.08

0.00

0.00

0.00

0.53

0.23

0.00

1.00

0.00

0.00

0.02

0.79

0.00

0.00

0.00

0.10

0.00

0.00

0.00

Seed 3 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

yH

+ [ ]

0.50

0.97

0.76

0.00

0.00

0.00

0.07

0.00

0.00

0.74

0.00

0.12

0.00

0.75

0.03

0.00

0.00

0.00

0.57

0.00

1.00

0.00

0.00

1.00

0.00

1.00

0.00

0.00

0.00

0.60

0.00

0.00

0.23

0.43

0.00

0.00

0.30

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.33

0.20

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.30

0.03

0.01

0.00

0.00

0.00

0.63

0.00

0.00

0.26

0.00

0.88

0.00

0.25

0.04

Seed 4 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

5

4

3

2

1

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

Seed 5 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Task Index

5

4

3

2

1

yH

+ [ ]

0.00

0.00

0.04

0.00

0.00

1.00

0.50

0.00

0.00

0.45

1.00

0.00

0.00

0.53

0.20

0.00

0.00

0.67

0.00

0.00

0.00

0.50

0.00

1.00

0.02

0.00

0.10

1.00

0.16

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.89

0.00

0.00

0.80

1.00

0.01

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.11

0.00

0.00

0.00

0.16

0.00

0.00

0.99

0.29

0.00

0.00

0.00

0.00

0.00

0.00

0.41

0.00

0.00

0.00

0.15

0.00

Seed 6 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Task Index

5

4

3

2

1

0.14

0.24

0.28

0.00

0.00

0.00

0.00

0.00

0.81

0.00

0.11

0.00

0.00

0.08

0.39

0.70

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.96

0.00

0.11

0.00

0.09

0.00

0.23

0.00

0.00

0.00

1.00

0.00

0.19

1.00

0.30

0.04

1.00

0.68

0.04

0.07

0.76

0.49

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.59

0.00

0.00

0.14

0.57

0.00

0.00

0.00

0.00

1.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

Seed 7 at 250M steps

(a) SDVT-LW

(

K

= 5

, α

c

= 0

.

5

, α

o

= 0

.

0)

.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

9 10

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00

0.00

0.00

0.10

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.05

0.00

0.00

0.04

0.00

0.00

0.00

0.30

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.07

0.00

0.00

0.17

0.00

0.00

0.05

0.00

0.00

0.00

0.15

0.26

0.29

0.00

0.00

0.01

0.16

0.01

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.11

0.00

0.00

0.00

0.00

0.00

0.00

0.00

1.00

0.09

0.15

0.00

0.00

0.00

0.02

0.00

0.00

0.13

0.01

0.10

0.04

0.35

0.00

0.00

0.00

0.07

0.00

0.00

0.00

0.19

0.31

0.35

0.00

0.00

0.01

0.20

0.01

0.00

0.00

0.00

0.13

0.30

0.00

0.50

0.19

0.00

0.01

0.00

0.48

0.28

0.17

0.04

0.40

0.00

0.00

0.08

0.00

0.00

0.00

0.20

0.32

0.33

0.00

0.00

0.01

0.21

0.02

0.00

0.00

0.91

0.39

0.00

0.00

0.00

0.05

0.00

0.00

0.87

0.01

0.20

0.05

0.52

0.00

0.00

0.00

0.13

0.30

0.00

0.50

0.19

0.00

0.02

0.00

0.50

0.26

0.17

0.04

0.39

Seed 0 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

9 10

8

7

6

5

4

3

2

1

0.00

0.00

0.00

0.12

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.27

0.32

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.02

0.00

0.12

0.00

0.00

0.00

0.24

0.00

0.00

0.11

0.23

0.07

0.14

0.04

0.09

0.11

0.39

0.05

0.22

0.00

0.16

0.00

0.00

1.00

0.06

0.07

0.04

0.05

0.03

0.05

0.38

0.00

0.09

0.10

0.00

0.00

0.00

0.26

0.00

0.21

0.00

0.27

0.20

0.00

0.22

0.01

0.00

0.00

0.03

0.00

0.01

0.36

0.00

0.00

0.13

0.30

0.08

0.16

0.12

0.11

0.12

0.54

0.05

0.29

0.73

0.17

0.26

0.00

0.00

0.04

0.08

0.01

0.02

0.12

0.05

0.28

0.03

0.44

0.16

0.00

0.00

0.00

0.31

0.00

0.21

0.00

0.25

0.20

0.00

0.23

0.02

0.00

0.00

0.02

0.00

0.00

0.00

0.31

0.00

0.21

0.00

0.26

0.21

0.00

0.22

0.01

0.00

0.00

0.02

0.00

0.35

0.14

0.00

0.00

0.03

0.33

0.02

0.03

0.68

0.03

0.04

0.05

0.25

0.16

Seed 1 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

9 10

8

7

6

5

4

3

2

1

yH

+ [ ]

0.21

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.33

0.00

0.32

0.00

0.00

0.00

0.00

0.32

0.01

0.00

0.00

0.00

0.00

0.00

0.00

0.34

0.00

0.34

0.00

0.00

0.19

0.00

0.34

0.23

0.00

0.00

0.22

0.05

0.00

0.51

0.00

0.00

0.00

0.78

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.06

0.00

0.00

0.00

1.00

0.00

0.00

0.19

0.00

0.00

0.00

0.00

0.00

0.01

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.74

0.38

0.00

0.00

0.00

0.00

0.00

0.00

0.43

0.13

0.12

0.00

0.00

0.00

0.03

0.01

0.09

0.00

0.00

0.03

0.00

0.21

0.00

0.25

0.40

0.00

0.00

0.00

0.00

0.00

0.06

0.34

0.08

0.00

0.24

0.00

0.59

0.26

0.00

0.11

0.00

0.33

0.00

0.33

0.00

0.00

0.00

0.90

0.33

0.41

0.00

0.01

0.52

0.00

0.36

0.37

0.00

0.00

0.00

0.18

0.00

0.00

0.02

0.00

0.06

0.00

0.19

0.00

Seed 2 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

9 10

8

7

6

5

4

3

2

1

0.00

0.00

0.00

0.00

0.00

0.21

0.00

0.00

0.00

0.00

0.17

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.44

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.63

0.01

0.00

0.26

0.00

0.35

0.11

0.00

0.00

0.00

0.14

0.00

0.28

0.15

0.00

0.00

0.00

0.17

0.00

0.00

0.49

0.10

0.00

0.00

0.00

0.15

0.00

0.31

0.15

0.00

0.00

0.00

0.40

0.00

0.00

0.00

0.00

0.00

1.00

0.39

0.00

0.10

0.00

0.00

0.40

0.00

0.00

0.00

0.00

0.00

0.00

0.25

0.14

0.00

0.00

0.28

0.00

0.00

0.27

0.00

0.00

0.26

0.00

0.13

0.81

0.04

0.00

0.00

0.00

0.40

0.00

0.00

0.00

0.00

0.42

0.00

0.00

0.00

0.01

0.00

0.00

0.27

0.16

0.00

0.00

0.29

0.00

0.00

0.28

0.00

0.00

0.29

0.00

0.15

0.15

0.12

0.26

0.02

0.00

0.00

0.11

0.00

0.41

0.16

0.00

0.36

0.16

0.44

0.25

0.04

0.00

0.01

0.23

0.00

0.00

0.04

0.90

0.00

0.00

0.00

0.01

0.28

0.00

0.21

Seed 3 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

9 10

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00

0.00

0.02

0.00

0.00

0.00

0.39

0.00

0.00

0.00

0.00

0.00

0.69

0.00

0.00

0.00

0.04

0.26

0.00

0.00

0.00

0.15

0.00

0.00

0.00

0.00

0.01

0.26

0.01

0.02

0.00

0.00

0.00

0.27

0.00

0.01

0.03

0.14

0.28

0.13

0.00

0.00

0.00

0.02

0.00

0.00

0.33

0.25

0.09

0.00

0.00

0.12

0.02

0.01

0.19

0.00

0.29

0.02

0.10

0.34

0.23

0.00

0.00

0.00

1.00

0.00

0.00

0.12

0.00

0.00

0.00

0.00

0.00

0.07

0.00

0.00

0.31

0.23

0.13

0.00

0.00

0.13

0.04

0.02

0.24

0.00

0.35

0.01

0.12

0.31

0.00

0.31

0.24

0.13

0.00

0.00

0.14

0.04

0.02

0.25

0.00

0.36

0.02

0.12

0.33

0.00

0.00

0.00

0.31

0.00

0.03

0.04

0.17

0.34

0.18

0.01

0.00

0.00

0.03

0.00

0.34

0.00

0.00

0.03

0.00

0.49

0.00

0.23

0.17

0.00

0.50

0.00

0.00

0.27

0.00

0.43

0.00

0.00

0.03

0.00

0.47

0.00

0.23

0.17

0.00

0.49

0.00

0.00

0.26

0.00

Seed 4 at 250M steps

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

9 10

8

7

6

5

4

3

2

1

0.00

0.00

0.00

0.08

0.00

0.00

0.00

0.00

0.07

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.17

0.00

0.21

0.00

0.00

0.18

0.00

0.22

0.01

0.00

0.00

0.19

0.00

0.00

0.00

0.22

0.00

0.27

0.00

0.00

0.25

0.00

0.27

0.01

0.00

0.00

0.26

0.00

0.25

0.00

0.00

0.49

0.20

0.00

0.26

0.01

0.00

0.21

0.00

0.00

0.01

0.12

0.00

0.27

0.00

0.00

0.51

0.20

0.00

0.30

0.01

0.00

0.23

0.00

0.00

0.01

0.14

0.09

0.28

0.09

0.04

0.00

0.07

0.01

0.35

0.05

0.04

0.05

0.16

0.00

0.27

0.16

0.41

0.12

0.05

0.15

0.00

0.03

0.00

0.06

0.14

0.37

0.01

0.29

0.00

0.21

0.07

0.00

0.00

0.13

0.13

0.00

0.00

0.61

0.00

0.10

0.08

0.00

0.14

0.00

0.11

0.00

0.10

0.00

0.69

0.05

0.00

0.00

0.38

0.00

0.04

0.03

0.00

0.09

1.00

0.23

0.01

0.39

0.08

0.04

0.16

0.00

0.02

0.00

0.04

0.15

0.47

0.00

0.30

0.00

0.18

0.05

Seed 5 at 250M steps

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

9 10

8

7

6

5

4

3

2

1

yH

+ [ ]

0.04

0.00

0.00

0.00

0.00

0.53

0.00

0.00

0.00

0.00

0.27

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.60

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.22

0.10

0.00

0.00

0.00

0.00

0.00

0.23

0.47

0.00

0.24

0.00

0.06

0.00

0.00

0.11

0.33

0.00

0.00

0.00

0.20

0.00

0.00

0.00

0.00

0.13

0.34

0.11

0.36

0.42

0.14

0.02

0.22

0.00

0.19

0.04

0.00

0.03

0.00

0.34

0.05

0.11

0.55

0.00

0.23

0.00

0.00

0.00

0.40

0.09

0.00

0.48

0.08

0.00

0.10

0.00

0.00

0.00

0.00

0.31

0.04

0.00

0.00

0.00

0.18

0.00

0.50

0.34

0.00

0.28

0.01

0.00

0.01

0.00

0.00

0.26

0.11

0.00

0.00

0.00

0.00

0.01

0.31

0.53

0.00

0.33

0.00

0.07

0.00

0.00

0.04

0.01

0.78

0.00

0.00

0.19

0.00

0.00

0.00

0.00

0.01

0.01

0.01

0.00

0.00

0.19

0.43

0.00

0.00

0.00

0.57

0.00

0.01

0.00

0.01

0.23

0.54

0.18

0.64

Seed 6 at 250M steps

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

9 10

8

7

6

5

4

3

2

1

0.06

0.02

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.03

0.00

0.00

0.23

0.06

0.00

0.18

0.20

0.22

0.05

0.00

0.00

0.07

0.00

0.05

0.26

0.00

0.00

0.25

0.23

0.00

0.02

0.29

0.25

0.13

0.19

0.06

0.35

0.00

0.48

0.03

0.00

0.00

0.24

0.06

0.00

0.44

0.19

0.25

0.13

0.00

0.23

0.25

0.00

0.00

0.01

0.50

0.00

0.00

0.03

0.45

0.00

0.00

0.00

0.00

0.04

0.01

0.00

0.00

0.04

0.35

0.00

0.00

0.28

0.35

0.02

0.01

0.31

0.27

0.64

0.38

0.08

0.32

0.57

0.35

0.00

0.00

0.00

0.00

0.00

0.00

0.35

0.00

0.00

0.04

0.00

0.57

0.01

0.10

0.02

0.14

0.00

0.61

0.00

0.10

0.03

0.00

0.00

0.00

0.00

0.13

0.02

0.00

0.00

0.00

0.02

0.50

0.00

0.00

0.04

0.47

0.00

0.00

0.00

0.00

0.06

0.01

0.00

0.27

0.04

0.16

0.00

0.39

0.00

0.13

0.03

0.00

0.00

0.00

0.00

0.20

0.03

0.00

Seed 7 at 250M steps

(b) SDVT

(

K

= 10

, α

c

= 1

.

0

, α

o

= 1

.

0)

.

E Ablation Results

In order to confirm the source of the gain obtained from the proposed method, we conduct an extensive ablation study based on SDVT and SDVT-LW on ML-10.

E.1 Occupancy Regularization

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

10

9

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.11 0.00 0.00 0.00 0.13 0.00 0.00 0.43 0.00 0.00 0.31 0.00 0.00

0.90 0.20 0.00 0.00 0.00 0.00 0.01 0.00 0.01 0.00 0.00 0.00 0.00 0.14 0.00

0.00 0.00 0.00 0.66 0.00 0.69 0.00 0.00 0.00 0.00 0.67 0.61 0.00 0.00 0.76

0.00 0.30 0.03 0.00 0.00 0.00 0.00 0.90 0.76 0.00 0.00 0.00 0.00 0.43 0.00

0.00 0.00 0.00 0.34 1.00 0.31 0.00 0.10 0.24 0.00 0.33 0.39 0.00 0.00 0.24

0.00 0.00 0.16 0.00 0.00 0.00 0.42 0.00 0.00 0.56 0.00 0.00 0.69 0.00 0.00

0.09 0.50 0.69 0.00 0.00 0.00 0.44 0.00 0.00 0.01 0.00 0.00 0.00 0.44 0.00

Seed 0 at 250M steps

(a) SDVT

(

α

o

= 5

.

0)

.

1. Reach

2. Push

3. Pick-place

4. Door-open

5. Drawer-close

6. Button-press

7. Peg-insert-side

8. Window-open

9. Sweep

10. Basketball

11. Drawer-open

12. Door-close

13. Shelf-place

14. Sweep-into

15. Lever-pull

10

9

8

7

6

5

4

3

2

1

yH

+ [ ]

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00

0.00 0.02 0.00 0.00 0.00 0.00 0.00 0.00 0.01 0.00 0.00 0.00 0.00 0.01 0.00

0.00 0.07 0.00 0.00 0.00 0.00 0.00 0.00 0.02 0.16 0.01 0.35 0.00 0.11 0.15

0.00 0.05 0.17 0.00 0.00 0.00 0.12 0.00 0.02 0.51 0.01 0.55 0.00 0.10 0.17

1.00 0.42 0.00 1.00 0.00 1.00 0.00 0.00 0.35 0.06 0.96 0.07 0.00 0.07 0.04

0.00 0.44 0.83 0.00 1.00 0.00 0.88 1.00 0.60 0.27 0.02 0.03 1.00 0.71 0.64

Seed 0 at 250M steps

(b) SDVT

(

α

o

= 10

.

0)

.

Figures 4b and 12 indicate that our occupancy regularization effectively constrains the use of higher indices as intended. We fix the default choice of α o = 1 . 0 that we found to work well on ML-10. Our default SDVT with α o = 1 . 0 scores a test success rate of 32.8%, which scored the best among α o ∈ { 0 . 0 , 1 . 0 , 5 . 0 , 10 . 0 } that scored 22.8%, 25.5%, and 21.0%, respectively. Also, our default SD with α o = 1 . 0 scores a test success rate of 30.8%, which scored the best among α o ∈ { 0 . 0 , 1 . 0 , 5 . 0 , 10 . 0 } that scored 22.8%, 28.4%, and 11.2%, respectively.

E.2 Dropout and Dispersion

Table 10 demonstrates the significant roles played by both dropout and dispersion in generating extrapolated dynamics during virtual training, thereby enhancing the test success rate. As outlined in Appendix E.1 of the LDM paper [35], virtual training without dropout does not exhibit a significant improvement due to decoder overfitting on the state and action inputs of the training tasks. In fact, virtual training without dropout can even lead to a decrease in test performance.

However, it is important to note that these factors are not the primary contributors to the empirical gains of our method. Dropout is specifically essential for virtual training, but not necessarily for other methods that do not employ virtual training, such as VariBAD and RL2 (as discussed in Appendix E.2 of the LDM paper). We verified this in our ML-10 experiment using VariBAD with dropout. The inclusion of dropout marginally improves the training success rate of vanilla VariBAD from 58.2% to 63.0% and the test success rate from 14.1% to 16.5%. However, the performance enhancement achieved through dropout is not as substantial for VariBAD as it is for SDVT-LW's virtual training.

E.3 Decomposition Distribution

Table 10 emphasizes the significance of hyperparameters such as subtask dimension K and categorical coefficient α c in determining subtask compositions. It is crucial to carefully set these hyperparameters based on the number of training tasks and their correlations. In our lightweight (LW) approach, we have selected K = 5 and α c = 0 . 5 , which effectively distributes subtasks across the ML-10 tasks. When α c is set to a small value, task classification collapses into a few subtasks, as depicted in Figure 4c. Conversely, with a large value of α c , the entropy of the composition is maximized, and all tasks exhibit a uniform probability distribution, as illustrated in Figure 4d. However, it is noteworthy that SDVT-LW outperforms VariBAD and LDM in test success rate, regardless of the specific values chosen for K , such as K = 3 , 7 , 10 . To alleviate the burden of hyperparameter tuning in our LW methods, we introduce the occupancy regularization.

E.4 Number of parameters

Table 11 indicates that SDVT-LW employs 3 and 2 times more parameters compared to VariBAD and LDM, respectively. The encoder of our method has more parameters than VariBAD due to the added categorical layer. The decoder of SDVT-LW includes the parameters of the dispersion layers. However, we find that the model's improvement is not mainly attributed to the increased capacity. Generally, a larger capacity does not necessarily guarantee better generalization.

VariBAD (hidden × 2) While selecting encoder (GRU) and policy hidden sizes, we experimented with 128, 256, and 512, discovering that 256 works best for all task-inference methods (LDM, SD only, and SDVT). Notably, VariBAD (hidden × 2 ) with hidden dimensions of 512 possesses more parameters than SDVT but exhibits an even lower test success rate than the vanilla VariBAD.

VariBAD (matched) We add MLP layers with hidden sizes of [1600, 256] into the encoder and [128] into the policy. The decoder's hidden size is increased from [64, 64, 32] to [128, 256, 160] to match the capacities of VariBAD and SDVT-LW. The components of VariBAD possess slightly more parameters than SDVT-LW. LDM employs two VariBAD encoders. While the matched baselines improve, they still lag behind ours.

E.5 Conditioning policy on belief

We employ VariBAD's stop gradient architecture to avoid the instability that may arise from the concurrent training of policy and ELBO objectives. This may lead to a potential vulnerability where the policy neglects contexts. However, in Meta-World, identical observations could belong to different tasks, such as 'Reach' and 'Pick-place.' Therefore, the agent must rely on the inferred context. To confirm this, we evaluate SDVT by masking either ( ω ϕ y ) or ( µ ϕ z , σ ϕ z ) in the policy input.

Our findings reveal that masking the contexts severely damages the training and test success rates. Especially, masking the Gaussian parameters of z is more critical than masking the categorical parameter of y . However, this does not suggest that y is redundant, given that the continuous context z is trained to include y 's information. Interestingly, the test success rate, with all contexts masked (i.e., solely relying on the current state s t ), achieves a success rate of 24.5%, surpassing all baseline scores. Although the policy does not directly employ the context learned by the GMVAE, it can still be effectively trained in the imaginary task generated by the learned GMV AE to prepare for tests with unseen compositions of subtasks.

E.6 Smaller number of rollouts

We use n roll = 10 , following the original setup of the Meta-World benchmark [69, 71]. Notably, this benchmark utilizes dense rewards, therefore the context converges rather quickly, even within the first rollout (Figure 2b). As a result, the performance across rollouts does not exhibit substantial improvement. To provide further insight, we report the success rates for smaller n roll in Table 13.

Peer reviews:
comments: SUMMARY:
This paper proposes SDVT, which addresses the non-parametric task variability problem in Meta-RL. SDVT utilizes a bi-level task encoding structure, namely a categorical encoding to address non-parametric variability and a VariBad-style continuous task belief. SDVT also utilizes imagined rewards to improve generalization. SDVT outperforms baseline Meta-RL algorithms on the hard ML10 and ML45 benchmarks.

STRENGTHS:
1. Dealing with multi-modal and diverse task distributions is an important problem in Meta-RL.
2. Presentation is clear, and the paper is easy to follow.
3. The methodology is sensible.
4. Experiment results are solid and convincing.

WEAKNESSES:
Overall the paper is clearly presented. I have a few concerns about why not use imagined dynamics as well as discussion with some more recent Meta-RL methods. Please refer to the Questions part for details.

QUESTIONS:
1. Why SDVT does not use imagined dynamics for training? Not using imagined dynamics may decease SDVT's generalization power.
2. Recent work [1] has shown very good performance on Meta-World ML10 and ML45. I'd like to see some discussions and comparisons between SDVT and the no-instruction version of [1]. Can transformers to some extend achieve the same ability to deal with non-parametric task variability? What's the pros and cons of SDVT compared to [1]?



[1] Bing, Zhenshan, et al. "Meta-Reinforcement Learning via Language Instructions." arXiv preprint arXiv:2209.04924 (2022).

LIMITATIONS:
The authors have discussed some limitations. I's like to know the computation burden of SDVT. Compared to baseline algorithms like PEARL, how computationally dense SDVT is?
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-23

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel context-based meta-reinforcement learning (meta-RL) methodology designed for tasks that can be decomposed into sub-tasks, such as the pick and place operation. To tackle this specific type of environment, the authors propose the adoption of a novel encoder variant capable of predicting a categorical latent variable in addition to the conventional latent vector, which follows a normal distribution. Moreover, the authors employ a technique called virtual training, where interaction with simulated environments is facilitated through the decoder, aimed at enhancing the algorithm's generalization capabilities. The authors conducted a comparative analysis of their proposed algorithm against state-of-the-art meta-RL baselines utilizing two meta-world benchmark datasets.\n\nSTRENGTHS:\n1. The new VAE architecture is well-motivated, well-formulated, and seems to have empirical benefits. These ideas might be helpful for researchers in the field.\n2. The evaluation is done on a well-known and very challenging benchmark. The authors did a fine job comparing to SOTA meta-rl works. \n3. The occupancy regularization seems like a very appealing solution to make the agent prioritize using lower indices first\n4. Overall, the paper is well-written and organized.\n\nWEAKNESSES:\n1. There lacks a formulation for the general problem definition, i.e. what does a \"non-parametric\" task mean? What do \"simple parametric variations\" (line 29) mean? I think the definitions in [2] might be the missing piece. This makes it hard to understand which type of environment this work is aimed at.\n2. Some design choices are not well motivated or compared to different design choices (see the questions section)\n3. There is no evaluation over simple and standard meta-rl environments (like simple goal-reaching). Although this is not the focus of this work, it is important to check the performances are on par with previous works.\n4. The experiment on the number of parameters in Appendix E.4 is very important, yet not complete. The authors show that even after matching the total number of parameters to SDVT, VariBAD performances are still inferior. This experiment does not prove that the gain in performance is not from the capacity of the encoder or decoder (which is very plausible since the length of the meta-episode is much longer than the one used in previous works). In order for the experiment to be complete, the authors should run LDM and VariBAD with an encoder, decoder, and policy of the same sizes as SDVT-LW.\n\nQUESTIONS:\n1. If I understand correctly, we can parametrize the tasks in meta-world, e.g. theta = (i, theta_l), where i is the task type number in meta-world and theta_l is the low dimensional parametric variation within each task type. So what do you mean by “non-parametric” tasks? \n2. In the occupancy regularization section, in the case K* is unknown, what is the motivation for making the effective K smaller? In your appendix you showed that the regularization factor of the occupancy regularization affects the success rate of SDVT, I wonder if it only affects the virtual training part. Can you please run the same ablation for SD as well?\n3. I couldn’t understand where the Gaussian mixture distribution is being modeled. You state (line 133) that the encoder is a multivariate Gaussian encoder, did you mean it is a Gaussian mixture encoder? If so, can you provide an exact formulation of this distribution? In case it’s not a mistake, and the encoder is indeed a multivariate Gaussian encoder, what is the input to the mean and variance layers of the encoder? The concatenation of h_t and y_t? \n4. Is the decoder dropout you present identical to the one in LDM? If so, can you refer to it in the paper?\n5. In your experiments, you used n_roll=10, did you observe an increased return between rolls? Can you provide a graph of the return of each episode within a test meta episode? Can you compare your method to VariBAD/LDM but with a lower n_roll (which is standard by previous methods)?\n6. In the imaginary reward generation section, you define the way the imaginary subtask compositions are being sampled (line 205), how is this different than LDM, can you compare the two? Why did you choose a Dirichlet distribution? Is there any theoretical standing for that? Please refer to [2], where the authors used KDE, shouldn’t it perform better?\n7. Why did you decide to keep only the imaginary subtask compositions constant during virtual training and not the final imaginary context (as done in LDM and in [2])?  \n8. In section 2.1 you formulate the problem for finite train and test tasks spaces, and in section 5.1 you state ML-10 and ML-45 have N_train=10 and N_train=45 respectively. This is a bit confusing because, to the best of my knowledge, these are actually infinite train and test task spaces (due to the parametric variations), can you clarify this?\n9. Related work - while mentioning MAML, I think it’s variant MMAML [1], might be more relevant to your case. Also, I think, [2] and [3] are also relevant to your work, please consider referring to them in your related work.\n\n[1] Vuorio et al. - Multimodal Model-Agnostic Meta-Learning via Task-Aware Modulation\n\n[2] Rimon et al. - Meta Reinforcement Learning with Finite Training Tasks -- a Density Estimation Approach\n\n[3] Marzari et al. - Towards Hierarchical Task Decomposition using Deep Reinforcement Learning for Pick and Place Subtasks\n\nLIMITATIONS:\nThe authors addressed the limitations of the paper", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an encoder architecture for meta-RL for tasks with non-parametric variation. The motivation for proposing a new encoder is that in order to achieve better generalization to a test-set of tasks with previously unseen non-parametric task variations, some specialized inductive bias can help. The proposed encoder is a Gaussian mixture model, where the idea is that the mixture components should correspond to the different subtasks. At training time, new combinations of the subtasks can then be created by sampling from the mixture model. The proposed method performs well on the meta-world benchmark.\n\n## Acknowledgment\n\nI have read the rebuttal and the following discussion and updated my review accordingly. The authors clarified aspects of the paper that were unclear to me. To reflect my new understanding of the paper, I raised my score from reject to weak accept.\n\nSTRENGTHS:\n## Originality\nThe proposed architectural change over the closest related method (LDM from Lee and Chung, 2021.) is straightforward but well motivated for the case of non-parametric task variation. Therefore, I would argue that the contribution is sufficiently original.\n\n## Quality\nThe motivation for the method is good and the method achieves strong empirical results on a relevant benchmark of non-parametric task variation. The basic elements of quality are therefore in place.\n\n## Clarity\nThe paper is clearly written.\n\n## Significance\nGeneralization is an important topic in the RL community. This paper proposes an architecture that seems to produce a significant improvement in generalization on task distributions with non-parametric task variation. Therefore, this paper could be an important contribution to the generalization in RL and meta-RL communities.\n\nWEAKNESSES:\n## Presenting the contributions\nThe scope of the contributions of the paper is not described accurately enough. The proposed method reduces to a new architecture for the encoder of the LDM (Lee and Chung, 2021.) method. This is only stated in the experimental section. In my opinion, the simple contribution is enough for a good paper when it results in a strong empirical performance, as seems to be the case here. But right now it is too hard to determine what is novel. I would recommend adding a paragraph in the introduction that more carefully outlines the specific contribution of the paper.\n\n## Justification of the occupancy regularization\n- The method section says occupancy regularization is crucial, but doesn't explain why.\n- The occupancy regularizer seems to act contrary to prior from the ELBO: the former pushes the $y_t$ toward zeros and latter pushes it toward a uniform distribution. You seem to want a $K$-hot prior parametrization, but I don't see why pushing specific dimensions of $y_t$ toward zero is needed.\n- In the imaginary reward generation, you use the empirical mean over the $y$ as the parameter of the reward distribution anyway, so I don't see why the dimension regularization is needed.\n- **Occupancy regularization** paragraph in the experiments section shows that the occupancy regularization does limit the use of the higher dimensions of the subtask decomposition, but none of the experiments show evidence that this is useful.\n- At the minimum, the occupancy regularization should be empirically validated. If it is actually useful, then it would be optimally incorporated as a part of the prior distribution for the VAE as that seems to be its role. Right now it is kind of an ad hoc regularizer that is not motivated strongly.\n\n## Missing ablations\n- The virtual training (SD vs SDVT) ablation is good, it clearly shows the benefit from virtual training. However, the other design choices are not ablated. Specifically the effects of the occupancy regularization, decoder dropout, and latent dispersion are not demonstrated empirically.\n- It is not clear why the -LW baseline is introduced. It was never motivated in the text.\n\n## Weak analysis\n- The **learned subtask compositions** paragraph has good kind of analysis about the results but figure 4 is not the easiest to interpret. Mapping the arbitrary task indices between the text and the figure to some similarity between the tasks is a challenge. Maybe you could cherry-pick interesting groupings of tasks, display them as bargraphs together, with some simple explanation of what is similar about those tasks? Then delegate the matrix of all task compositions to the appendix. Or use some other visualization of the subtask decomposition. The point is the chosen presentation is not easy to interpret.\n- **Occupancy regularization** shows that the occupancy regularization does limit the use of the higher dimensions of the subtask decomposition, but none of the experiments show evidence that this is useful.\n- **Generated imaginary tasks** paragraph has a good idea, but the figure is not very informative. You can see that the policy tends to do different things for the different decomposition vectors, but from the figure it is impossible to tell whether these behaviors are usefully and meaningfully different. I would recommend identifying some of the specific behaviors the different dimensions cause and then show how those can be controlled by choosing different subtask vectors.\n\n## Minor issues\n- The syntax for the dropout is potentially confusing, as it is applying the dropout directly on the states and actions. Unless $s_j$ and $a_j$ are the embeddings already, but that is equally confusing.\n- Would be good to explain what is the target $h_t$ in the latent dispersion objective in the text.\n- Figure 2 a and 3 a are confusingly similar. Not clear why you need to show both.\n- The standard PEARL isn't really a zero-shot meta-RL method because it requires full trajectories to update the belief.\n\n## Summary\nIn my opinion the experimental results are good enough for a strong paper. However, the motivation of the method and presentation are not good enough for publication in its current form. Most importantly, it is too hard to understand the specific contribution of the paper over the related work and the paper lacks ablation studies on central components of the algorithm. I would be willing to raise my score if the issues listed above are addressed in a revised version.\n\nQUESTIONS:\n- Why is the occupancy regularization needed? Couldn't you just pick a large $K$ and accept that some of the dimensions are non-informative?\n- Why would overfitting lead specifically to low predicted rewards for the less often seen states? Why not high?\n- How does the proposed method handle tasks where the subtasks need to be performed in sequence?\n- How is LDM limited to parametric variation between the tasks only?\n- What is the effect of the decoder dropout on the performance of SDVT?\n- What is the effect of the latent dispersion on the performance of SDVT?\n- What is the effect of the occupancy regularization on the performance of SDVT?\n\nLIMITATIONS:\nLimitations of the proposed method are not discussed at a sufficient level of detail.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nMeta-RL has proven successful in generalizing RL across tasks. Current methods struggle to generalize across tasks with novel sub-task decompositions. The authors introduce SVDT, a GM-VAE method that decomposes tasks into sub-tasks that help with this form of generalization. Virtual training is employed to further increase robustness to unseen sub-task compositions. Dropout, regularization, and dispersion is used to further enhance performance. On Meta-World SVDT significantly outperforms existing methods.\n\nSTRENGTHS:\nThis paper has many strengths which I list below:\n\n- It is very well written and clear\n- Strong consistent performance (8 seeds) with respect well known baselines\n- Combines strengths of related works into singular framework\n- Adapts virtual training to the sub-task setting\n- Strong ablations and qualitative analysis (in both main paper and appendix)\n- Code is provided\n- Website is very informative\n\nWEAKNESSES:\nThe main weakness is that the approach has many hyperparameters and a lack of analysis as to why SVDT-LW outperforms SVDT with respect to generalization. In line 295, the authors have a hypothesis. It would be great if they could include experiments to validate this hypothesis in the paper.\n\nThe authors could also expand the offline skill learning related works section to include relevant sub-task decomposition works such as [1,2].\n\n[1] Salter, Sasha, et al. \"Mo2: Model-based offline options.\" Conference on Lifelong Learning Agents. PMLR, 2022.\n\n[2] Ajay, Anurag, et al. \"Opal: Offline primitive discovery for accelerating offline reinforcement learning.\" arXiv preprint arXiv:2010.13611 (2020).\n\nIn addition, the authors often use the term non-parametric task generalization. The generalization settings the authors consider in their paper are often called modularity/composition-based generalization (see [3]). Might be worth using these terms as well.\n\n[3] Khetarpal, Khimya, et al. \"Towards continual reinforcement learning: A review and perspectives.\" Journal of Artificial Intelligence Research 75 (2022): 1401-1476.\n\nQUESTIONS:\nCould the authors clarify why tasks that can be decomposed into sub-tasks are not considered parametric? Can’t the subtask ordering and labels be the parameters?\n\nLIMITATIONS:\nSee other sections.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}