Parameterizing Non-Parametric Meta-Reinforcement Learning Tasks via Subtask Decomposition
Authors: Suyoung Lee
Myungsik Cho
Youngchul Sung
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep reinforcement learning
Meta-reinforcement learning
Subtask decomposition
TL;DR: We enhance the generalization capability of meta-reinforcement learning on tasks with non-parametric variability by decomposing the tasks into elementary subtasks and conducting virtual training.

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
Suyoung Lee, Myungsik Cho, Youngchul Sung'
School of Electrical Engineering KAIST
Daejeon 34141, Republic of Korea fsuyoung.1, 1, ms ms. cho, ycsung)
Abstract
Meta-reinforcemeet learning (meta-RL) techniques have demonstrated remarkable success in generalizing deep reinforcement learning across range of tasks. Never- theless, 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 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 com- mon subtasks. Additionally, we propose a virtual training procedure, specifically designed for non- con-parametric task variability, which generates hypothetical sub- task compositions, thereby enhancing generalization to previously unseen subtask compositions. Our method significantly improves performance on the Meta-We Meta-World ML-10 and ML-45 benchmarks, surpassing current state-of-the-arr techniques.
Introduction
Meta-reinforcemmenn learning (meta-RL) constitutes 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-RI. creates an inductive bias regarding the task dynamics and subsequently develops a policy based on this knowledge. Despite its significant contribution the generalization capability of traditional deep RL, meta-RL is susceptible test-time
distribution shifts, which restricts its applicability to familiar in-distribution test tasks [17 36 39
To tackle this limitation, recent out-of -istrf--istributiinn (OOD) meta-RL approaches have emphasized distinct training and test task distributions, thereby achieving enhanced performance on unseen OOD test tasks with interpolated the parameters of training and test tasks are drawn from disjoint distributions, these tasks remain qualitatively similar, as they be expressed in shared parametric form representing the task dynamics (e.g., the same "Pick-place" task with OOD goal positions in Figure lb].
slightly extrapolated training dynamics [12 36 41
35 I. Although
In this study, we explore more general meta-RI. framework that addresses non-parametric task variability (69 771, a topic that has received limited attention in prior research. In this context as of Figure Ic 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 parametric embedding applicable to various tasks 9 46 (4). Within a non non-paramttric framework, it may not be feasible to employ unified and generalizable parameterization of training tasks using standard meta- RL techniques.
"ooresponding author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Musc
Muran
Murai
Meest
(a) "Pick- "Pick-place" standard in- distribution mela-RI. setup.
(b) "Pick- plack"" out-of- distribution meta-RI. setup. tween "Pick-place" and "Sweep-into."
(c) Non- parametrie task variation be-
Problem Setup. Visualizing different meta-RL scenarios with Meta-World tasks 669 The circles and crosses represent the object and goal positions, respectively. Solid blue objects indicate training tasks, while empty brown objects indicate test tasks.
Figure I:
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-parametrie task variability in meta-RL, our primary strategy involves decomposing each paraamerettic task into set of shared elementary subtasks. We then parameterize each task based on the types of subtasks that constitute it. Despite the con-amametric task variability, tasks may share elementary subtasks. For instance in Figure Ic "Pick-place" task can be decomposed into subtasks: "grip object" and "place object," while a ""Seeepintoo" 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 (I) 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 Gaussian mixture variational autoencoder (GMVAE) [29] [5) Our GMVAE encodes the trajectory up 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 -anaaaraettric variations. Consequently, the policy, using the learned subtask composition, can readily generalize new tasks comprising previously encountered subtasks. To further enhance generalization to unseen compositions of familiar subtasks, we propose virtual training (VT) process 35 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- Meta-World World ML-10 and ML-45 benchmarks 711,, widely used meta-RL benchmarks comprising diverse non-paramttric 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 31.2% on MI.-45, which improves the previous state-of-the oatte-of-the-rr 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...,,,
is defined by tuple comprising a set of states S, set of actions A. a reward function R(r++1 St, a, St-1), ). a transition function T(s-1133 4z), an initial state distribution To(ss), a discount factor y, and 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 M(Main)) over a set MDPs Mirain Each MDP Mr (S, A. R2. Tk, To,k, 7, H) is defined with unique reward function Rg, transition function Tx. and initial state distribution To Unlike in a multi-task setup, the agent in the meta-RL. setup does not have access to the task index that determines the MDP dynamics. The training objective is optimize the policy g with

---

parameters vo to maximize the expected return across all MDPs: maxy
(v))
where (v) Eo. Ne. Ei=o' 7 (T(41l ot, st+1) During meta-testing, standard in-
distribution meta-RL RL methods are evaluated on tasks sampled from the same distribution p(Msss)
as the training tasks, i.e., Mirain Miest .M (Figure Ia). In contrast, OOD meta-RL methods assume strictly disjoint training and test task sets, i.e., M Mimin Mies and Murann O Miea 0 (Figure lbp.
2.2 Bayes-adaptive Rent-Reeiforreement Learning
Since the true task index 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 Bayes- adaptive agent [38 10 [6] achieves this balance by updating its belief (R I') about the MDP based on its experience Ts {so. ap, r1, $1, d1 r2...... St dt-1, Tx, st). The agent's belief over the MDP dynamics at time t can be represented as posterior given the trajectory, i.e., by T) p( R, T..:)). By augmenting the state with the belief to form a hyper-state space S S B, where B is the set of belief, 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 folll rollout episodes (i.e., H Aooll H steps) of the same MDP:
Jpai(v) = Ebo, TI, mm
Σ
? R (ret2lsf at stx)
where T (s 1tt (r, Ea, [T(s4 (1, a:)] [(b241 = PR, T(444) is the transition dynam- ies and R+ R+ (re+1l8f at, 8(41) Re,,, [R(re+1 18t, at, 81+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 [7] solves the inference and posterior update of the belief by combining meta-RL and approximate variational inference (28). At each timestep t, recurrent encoder qs, encodes the expe- rience T. into a hidden embedding h = qd, (T=). The approximate posterior belief over the dynamics can be represented the parameters of multivariate Gaussian distribution: b. (Juz (he), To. (h+)). where Ma= (-) and oo_(1 (.) are neural networks. The latent context 2t -N (he), o (he))
is used to estimate the MDP dynamics: pen (T $j, aj. Sy+1, 21) and po, (s |s.. a, 21) for j = O..... H 1, including the past and future. Then the problem of computing the poste- rior over the dynamics p( RR T)xx) reduces to inferring the posterior qa (2x (T.)T.)), where ó f $z).
A separate policy network Ty (ax (a, b,) is trained to optimize the BAMDP objective in Eq. 11.
2.3
Task Variability
The term "non-pane parametric" in the context of task variability is introduced in the Meta-World paper 699 is used to distinguish the task variability in WeradWorld manipulation tasks from more simplistic parametric variations exhibited in standard MuJoCo tasks, such variations in target goal positions, directions, and velocities. However, the term "non "non-paraae parametric" might lead to misunderstanding. Because it could suggest that there 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 parammarretrric terminology, we guide the readers to view the scope of our work through the lens of modularity or composition-bant generalization [27].
3 Method
In this section, we introduce our novel meta-RRL method, named Subtask Decomposition and Virtual Training (SDVT), to handle parameeriiic 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 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.
3

---

Pick-place, K= 5, dimlz) 5
Optional dispersion for SDVT (not used for SD)
4
R
5 0 8
Dim.
3 -
Decoder
0 H
4
Encoder
831
Pea
h
NV (Now (he,s)) w. 0. (he, ve) x
a(-1
+ Cat (wra (Wa))~ yc
8
p6r
8
14
8541
-1
Policy
<->
Reconstruction
logo,
-2.
Stop gradient
Dropout
-5.0
100
200
300
400
500
1000
2000
3000
4000
2000
(a) SDVT architecture.
(b) Task inference on "Pick-place."
Figure 2: SDVT architecture. (a) Our proposed architecture incorporates three main components:
the encoder, decoder, and policy, An online trajectory is encoded into categorical (y) and Gaussian (2) latent contexts. These contexts, which are trained to reconstruct the forward dynamics, are utilized by the policy network. This structure is also applied the virtual training, as illustrated in Figure [3a with an optional dispersion layer integrated. (b) An example of the learned task inference process within meta-episode (H+ 5000 steps) "Pick-place is shown. We report the values for each dimension of contexts: wo, (hx), Pá_ (he, Vb), and log σo. (he, Wt).
3.1 Subtask Decomposition (SD) with Gaussian Mixture Variational Autoencoder
Our goal is to meta-learn set of elementary subtasks and meta-leam the decomposition of each task into composition of these subtasks. The core of our method focuses on meta-learning the approximate subtask composition 1/+ E A sampled from K-class categorical distribution, where A denotes the K --dimensional probability simplex. For example with K = 3. we want the subtask composition y/+ be learned somewhat like (0.5, 0.5, 0.0) for "Pick- Prick-place" and (0.0, 0.5, 0.5) for "Sweep-into,"" where each dimension of y, represents the weight corresponding to subtasks in the order of "place object," object," and "push object." To capture such subtask information, we use 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- K-dimensional mixture proportion ye. resulting in each class representing 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 Gaussian distribution. The learned subtask composition y represents the agent's 's belief at time about the current task's 's decomposition into subtasks. It is crucial to note that we don't 't necessarily want Ye to be one-hot embedding representing the subtask that the agent is solving at time t. We want y/ to represent a mixing proportion [56] over all possible subtasks that the agent believes be relevant to the current task, including past and future in Figures [ and 4 This distinction is vital to the model's '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 based on the VAE of VariBAD 77),, consists of three main components; the encoder, decoder, and policy networks parameterized by ó {on. dy $z). (0_, n2 }r and v. respectively.
(1) The encoder is defined as (o 2, h4) = gs, (ve ho) (2 he, U1). A recurrent network en- codes the past trajectory Tir into a hidden embedding he = Gon. (Ta).. First, the categorical encoder Ja, (3e)/+)= Cat(wa, (ht)) samples yt. where ws (he) E A. We use the Gumbel-Softmax trick 23 with a high temperature (r 1) when sampling y/2 to form soft label. Then the multivariate
Gaussian encoder as. (z /ig, vt) .N(Hss (he, Vt), o3 (he, yt)) samples a continuous latent context 2,, which contains the parametric information of the subtasks, in addition to the categorical subtask information of yt-
(2) The decoder is defined as (o Vt. 2.) p(s)po, (2x (2xly) (elly))o9g.077 (T (zx)- reconstructs the reward and transition dynamics for all transitions in the meta-episode T. using the latent context 2, as in Eq. (4). We assume uniform prior of subtask composition p(y:) Uniform(1/ K) and Gaussian regularization pe, (zelge) NNee, (ye), 0z, (yc)). This encoder-decoder architecture

---

allows both the approximate posterior (V,,, 24 ha) and the prior p(z)) to follow Gaussian mixture distributions.
(3) The policy network, The (ax|*,, bt). is trained separately conditioned on the belief b, = (ne, (l4, Ut), Ve. (h,, yt)) that parameters of the Gaussian context. In practice, we also pro- vide the parameters of the categorical encoder wo, (ha) to the policy, which we find to improve the performance. The parameters of the distributions (w2, Hob To 140,, and oa_) 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 = O,..., ......H* H over the trajectory distribution d(Mr, T, induced by the policy in MDP M
(L 0) = Ed(Me: (th.)
(2)
JGMVAE = RRroe + aTroe + agJre + CioJeat-
(3)
In addition to the reconstruction objectives JR- and Trree we have additional regularization Jrg and categorical Jea objectives:
H+ Σ
JR-ree
Σ j=0
Pon non (r)t1ls), 85.1,,,,,513 S1+1,
Free =
log pe, (85+1183, 20)
(4)
j=0
po, (2+181)) 11))
p(y))
Jeg = log
Jea = log
(5)
90. (Zellte, y1)
qs (ye he)
The regularization objective Jreg minimizes the KL divergence between the learned posterior Gaussian distribution go. (2elhe,3) 36) and learned Gaussian priors pe. (zelye). Unlike the standard VAE that assumes a standard normal prior, we learn K distinct Gaussian priors conditioned on yt. The categorical objective Jear maximizes the conditional entropy of y given h, and prevents collapse. Refer to Appendix [ for the derivation of the EI BB objective. The reconstruction objectives in Eq. (4) are computed for all timesteps from the first to the terminal step the meta-episode. Therefore, the subtask composition ye at time i is not necessarily a one-hot label of the belief about the current subtask at time f, but 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. ①T 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:
H
(,, (MMap(MMm))
Jnh((())
ELBO.(,,0)
(6)
Occupancy regularization The dimension of yr or the number of underlying subtasks K is crucial hyperparameter that should be determined based on the number of training tasks Nurain 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 meta-learn the number of effective subtasks as well. First, we assume K* < Nurii otherwise each task will be classified into separate subtask with one-hot label, preventing learning shareable subtasks. We start with a sufficiently large K = Nirai and regularize the ELBO objective to progressively reduce the number of effective subtasks (non-zero components) occupied in y/ with the following occupancy regularization that penalizes the usage of larger indices in the subtask composition:
Joce 1--log log K (e-K+1 KK2
y/-
We calculate the dot product exponential weights and the subtask composition ye to penalize the occupancy of higher dimensions of yı. We scale the dot product by log K to match the scale to the upper bound of Jem- We add Jace multiplied by a coefficient a to the GMVAE objective in Eq. (3). Consequently, the agent prioritizes using lower indices in the decomposition subtasks and sparingly uses higher indices for rare subtasks as in Figure 46 and in Appendix
represent frequent
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 2.. When
5

---

Imaginary subtask composition
Decoder
03
pop
1
1++1
M (Me, M). 03, (he, 1) 3))--)
PoN
4,30
Encoder
N (Ho, (hr, Bt): oz, (Pee)
A1
52
0A
Maximize
(b) [1,0,
(C) - 10, 1,0,0 0, 0, 0,0]
q
Cat (wes (Ree)) ~
H Σ
X
Stop gradient
Po.. do.
Policy
623
Dropout
(a) Imaginary reward generation.
(d) 1 = 0, 0].
(c) = [0, 0,
Figure 3: Virtual training. (a) Generation of imaginary rewards with the decoder conditioned on a fixed imaginary subtask composition ÿ. (b)-(e) Examples depicting the diversity of generated imaginary tasks, where K and all states are from the Meta- World "Reach." " Red rods and blue circles represent the trajectories of the gripper and object, respectively. These trajectories, while generated by the same policy, differ across various imaginary subtask compositions y.
such overfitting happens, the latent context loses its task-informative value, leading 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 b untouched: Pex (r)+1 DO [sy, aj, 11+1), =4) and pe, (# DO [sy, ay], =t). This dropout application 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 П. but their generated tasks are limited to parametric variations, c.g., generating tasks with unseen goal positions. Such methods 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 to extend the process the non conpaarmettric setup by conditioning the decoder on an imaginary subtask composition y. At the beginning of each meta-episode, we randomly determine with probability Pe 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. [ we utilize the dispersion structure for our GMVAE to support extrapolated generalization of the latent context 2,. Instead of directly using 2n as the decoder input, we insert an additional MLP po, before the decoder to expand the dimension of the context (the dotted box in Figure 2ap. The MLP output he pe, (2t) is trained to reconstruct the embedding he by appending aaa = -xaal/he h. 112 to the total GMVAE objective in Eq. (3). We then the dispersed context h, instead of 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 the distribution of training subtasks but with unseen compositions of subtasks. At the beginning of a virtual meta-ep, we randomly sample an imaginary subtask composition ÿ Dirichlet(ÿ) fixed for that virtual meta-episode, where the concentration parameter ğ E A K is the empirical running mean of all y/ over training. By re- placing the real ye with the imaginary subtask composition J, sample imaginary context
2 (us. (he, V) 03, (he, ;)), imaginary dispersed context hy po, (žt), and finally the imagi-

---

nary reward f(+1 ~ Pon ( [DO [st, at, S141] he) accordingly using our GMVAE. We replace the reward for the next timestep, Tt+1, with fr+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, Ei 7f 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 SDVT- LW 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 VAE (.. 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- initialization of 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 augmented 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 -iitt-iiittiittoionn tasks.
methods [13] 54] 48
73 53 learn common
18
46 32
34
45
74
40 use a recurrent or memory-
-fitt--dist 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
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 [I] 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 parameerit task variants.
44 33]
Skill-based meta-RL Our task parameterization based on the subtask decomposition is related to the recently spotlighted skill-based meta-RL methods [1] [1 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 [4] 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 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
7

---

the skill decomposition 688] 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, 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 Metannearning 10 (ML-10) benchmark, consists of Nuaio = 10 training tasks and Nust = 5 held-out test tasks. We denote each task by 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 Nurin 45 training tasks and New 5 test tasks. Refer to the tables in Appendix [D. for the set and indices of tasks. The agent must maximize its return from experience while exploring to identify the initially unknown task dynamies within meta-episode of H+ 5000 steps that consists of Proll = 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 = Nuain and apply the occupancy regularization (a, = 1.0) with o₂ 1.0. We also evaluate lightweight (-LW) version of ours with smaller K = with a. = 0.5 and without the occupancy regularization (a, 0.0). To ensure fair comparison and to exclude the gains from orthogonal contributions, we compare SDVT with ofatc-of-the--e meta-RL methods that do not require any refined offline demonstrations, ensembles, extensive test-time training: RL 2 [9) MAML [3], PEARL [46. and VariBAD 774. 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 Gaussian mixture reduces to LDM, LDM without virtual training reduces to VariBAD, and VariBAD without a VAE decoder reduces to RL.. Our implementation is available at https://gi tww...th.....m com/ aooomm/uuung--/ 1ee/SDVT
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 predefined threshold. A rollout episode is considered successful if the agent ever succeeds at any timestep during the rollout episode. The success of meta-episode is defined the success of the last (10th) rollout episode. In Table [ we report the success rates of 750 (15 tasks x 50 parametrii variants) meta-episodes at 250M steps for ML-10 and 2500 (50 tasks 50 parametric variants) meta-episodes at 400M steps for ML-45 Likewise, we report the returns of the last rollout episodes.
Evaluation metric
5.2 Results
Table I lillustrates 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 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 training success rate. In particular, outperforming VariBAD underscores the limitations of employing single Gaussian distribution to model the latent task space in cases of non-parametric variability.
On the test tasks, SDYT and SDVT-LW substantially outperform all baselines, even outperforming LDM, which surpasses VariBAD with parametric virtual training. Our gain is attributed to 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 distinet subtasks "move gripper" and "tighten gripper" with larger K. In contrast, SDVT-LW scores higher
*Our aggregation method diverges from the atypical method employed in the Meta- Meta-World paper 771. which presents the main results as the average of the maximum success rate of each task (details in Appendix C.5
8

---

Table 1: Meta-World V2 success rates and returns. We report the final success rates (%) and returns of our methods and baselines averaged across training tasks and test tasks of the ML-10 10 and MI,-45 benchmarks. All results are reported as the mean success rate + 95% confidence interval of seeds. Individual scores of all tasks are reported in Appendix D.I
Success Rate
Return
MI.10
MI.-5
MI.10
MI.-45
Methods
Train
Test
Train
Test
Train
Test
Train
Test
SDVT
77.2++3.0
32.8+3399
55.6+44.2
28. 28.12
3656:162
1225=160
2379-=214
8394
SDVT-LW
62.11:4.1
33.41:.00
50,414.1
31.2.11.2
33741137
1527-214
2294-1202
894.127
SD
77.0115.9
30.817.7
61.0111.7
23.0.15.1
36301:241
1112-190
2672.1.9
786.169
SD-LW
75.5.:5.5
26.218.7
56.711..
25.4.1..9
3525
1043-234
2578.164
793.149
RL
67.4::4.4
15.1
58.0+0.4
11.8+3.2
1159+
715=33
1411=22
663+
MAML
42.2.14.55
3,9-3.7
32.011.4
19.816.3
1822.1136
439-78
1388-110
658.196
PEARL
23.2.11.9
0.8
10.311244
6.7.3.3
1081177
340155
5971112
506.1122
VariBAD
58.2.18.9
14.116.1
57.0
22.14:3.5
305511466
9191144
2492.1447
762:140
LDM
56.7=12.3
19.8=6.0
54.1#0.9
24.8#
2963+626
1166=264
2515=167
768-663
test success rates than SDVT, presumably due its coarser decomposition that allows imaginary compositions to encompass 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 substantial increase in rollouts and updates during tests, Detailed ablation results are reported in Appendix [. For additional results such as rendered videos, refer our webpage https:// http:///iiees.gogle... wtt..//o...///oi//////o/ooooeeaa sdvr//urittneeippss
of potential future work.
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 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 motivation, visualize the converged subtask compositions in Figure 4 We find that such converged subtask compositions shared by qualitatively similar tasks. For example, in Figure 4a tasks (3) "Pick-place," (7) "Peg- PPeg-insert-side; and (10) "Basketball," which require placing an object at a goal position, share subtask indices and 3. Additionally, simple tasks that require moving the gripper in straight line: (1) "Reach," (5) "Drawer-close." and (8) "Window-oper These observations suggest that 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 the training tasks. For instance, (3) "Pick-place" and (14) "Sweep-into" share subtask 3, but not subtasks I or 5, revealing the potential for virtual training to be effective.
are classified into the same subtask 2.
Occupancy Regularization Figure Soureeiidiia indicatss that it is important to set the subtask class dimension K and the categorical coefficient a, appropriately according to the number of training tasks and the correlations among them. Otherwise, the decomposition may suffer from collapse degeneration. With the occupancy regularization, we can avoid the extensive search for the optimal subtask dimension K and corresponding ae- In Figure 4b we find that 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 ÿ in Figure When set ÿ as a one-hot vector, the policy

---

0.000.29 wt0.00.000000000000
000.0000000. 0.000.11.0000000 1..011000000 1.000.001.00
000.1.00.0000000
.0000.00..200000
0.000.000.00
0.000.000.00
0.000.22 2.0000220000
0.000.000 www.040.000oooo
00..000000.00.000000
0.000.040.00
1.000.000.00
Yu Vw
0.000.48 0.0.0.0000.000000000000000000000
0.00000000000 0tt.000000000......
0..000000.000000000
000000000.000000 0ww.0000000000.0 00000000.73044
0.190.060.89
1.0000000.00
2
0..00030 0,000000000 0.000.00000
0.00000.00 0.0.0000.000.00 .ww.00000000000.
000.00
0.000.000.0
00
2
.00 0.01001.00
0.000.01 00...01..81
0.81 w00.000.000.0..o
0.000.000.00 000.900.000.00 000.130.000 .00
1.000.00000000
1.000.000.55
0.810.900.11
O.
1.000.000.00
1. R-ach
Push
3. Pick-place
Door sporapen
Drawer-closs
6. Button-press
T. Peg. waag-insertssso
9. Swweop
10. Basketball
11. Drawer-
12. Door-close
13. Shelf
Sweep-ints
15, Lever pull
3. Pick-place
Pog lnsiinnteettt side
10. Basketball
1. Reach
Drawer-close
8. Window-apen
Window
SDVT-LW (K
5,0 0.5, 0% 0.0).
Sced at 250M steps
1.00
0.0000000.... 10.00.00000.00000000 0.0.000000000000
0.00
100 o.
0. 07
0.000
0.000.000000000000 0.00000.1111.
0.00000000
0.0000.00 0.001000
1.000000 0.000.000000000000 0.00 0,00000
100 0.0000
00
0.30
0.00
700
0000
0.0000000000000
0.00 0.00 0.00......
10.00.0000 00.00....00 0.00 8.15 i. 6.2% 0.0.0..0000000
6.160.01 0.00
DO 9.000
0.00
00 009.100:00
0000
1:00
00 00 00 000
0.00
0.000.000 0.000 0.00
14
1.00 0.000.07 0.000..0000 0.00 0. 3. 0:35 0.000.00000000000 0.20001....000.. 0.200.01.0000
0.00 0.00 9: 1300 30.0 00 8.50 19 00 0.01 0.00
(c) SDVT-LW (K = 5, Cz. = 0.1 Cx = 0.0).
0.08
0.000000..000
0.000.00
0.00
0.00 0.000.3..133 0.30 0.300.000 00 0.50 0.19 00 0. 0.5000 0.200.6.66 0.040
0.110 0.39 0.000.0.0000
0..0000000000000..00
0.000000000000000000000
0.00 0.00 0.20...05
0.39
8.260 D. 100 2 2002
230
39
0 1
24 38
0
2
0
20
21
220
15
0
D.
22
190.2292222 22 20 20
200
19D 200.
18 0.21
Neach
2. Puee
3.
Sweepp
10 Tantett bxl
st. Drwoerepp spp
Sholrpacce
14. Sampphii
15 Leverpaa
190
22
190
200
25
0
21
0
0
200
20
10
14 15
(d) SDVT-LW (K = 5, az
2.0,
0.0).
(b) SDVT (K~ = 10, ac
1.0, ao 1.0).
Figure 4: Learned subtask compositions on ML-10. (a) Default SDVT-LW (b) Default SDVT. (c) and (d) SDVT-LW with varying Oge Each column displays the terminal subtask composition (yz+) of each task learned after 250M training steps of training, averaged across 50 parametric variants. The decompositions at different timesteps can be found Appendix D.3 The results shown are from the first random seed, as different seeds yield distinct decompositions. Results of other seeds are provided in Appendix D.4 D.4
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-RI. with non non-parametric task variability. This improvement is achieved by meta-learning shareable subtask decompositions and executing virtual training 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 RL to accommodate wider range of task variability is critical research topic. Incorporating orthogonal approaches such as using offline demonstrations or test time training techniques into our method could lead interesting future work addressing the limitations. We optimistic that study lays a robust foundation for future research in this field.
Acknowledgements
This work was supported in part by the Institute of Information & Communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) 2022-0-00469,
Development of Core Technologies for Task-oriented Reinforcement Learning for Commercialiiation of Autonomous Drones, 50%) and in part by the Institute of Information Communications Technol- ogy Planning & Evaluation (ITPP) grant funded by the Korea government (MSIT) (No. N02-002--0001 24,
Development of Artificial Intelligence Technology for Self-Improving Competency-Awe Learning Capabilities, 50%)
10

---

References
[1] A. Ajay, A. Gupta, D. Ghosh, Levine, and P. Agrawal. Distributionally adaptive meta reinforcement learning. In Neural Information Processing Systems (NewrIPS), 2022.
[2] 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.
[3] 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.
[4] 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-140 1425, 2020.
[5] 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.
[6] B. Chen, W. Deng, and H. Shen. Virtual class enhanced discriminative embedding learning. In Neural Information Processing Systems (NeurIPS), pages 31:1942-1952, 2018.
[7] 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.
[8] 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 Reprexentations (ICLR), 2017.
[9] Y. Duan, Schulman, X. Chen, P. L.. Bartlett, I. Sutskever, and P. Abbeel. RI.3: Fast re- inforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1611.02777 2016.
[10] M. O. Duff. Optimal Learning: Computational Procedures for Bayes-adaptive Markov Decision Processes. University of Massachusetts Amherst, 2002.
[11] B. Eysenbach, A. Gupta, Ibarz, and S. Levine. Diversity is all you need: Learning skills without reward function. In International Conference on Learning Representations (ICIR), 2019.
[12] R. Fakoor, P. Chaudhari, Soatto, and A. J. Smola. Meta-q-learning. International Conference Learning Representations (ICLR), 2020.
[13] C. Finn, P. Abbeel, and S. Levine. aggossag meta-leaming for fast adaptation of deep networks. In International Conference on Machine Learning (ICML), pages 70:1 1221126-1110 1135, 2017.
[14] C. Florensa, D. Held, X. Geng, and P. Abbeel. Automatic goal generation for reinforcement learning agents. In International Conference Machine Learning (ICML), pages 80:1515- 1528, 2018.
[15] Garage-Contributors. Garage: A toolkit for reproducible reinforcement learning research.
https://github. http://ww..aaoeoooooooooo//e/eee garage 2019.
[16] M. Ghavamzadeh, Mannor, Pineau, and A. Tamar. Bayesian reinforcement learning: A survey, Found. Trends Mach. Learn., pages 8(5-6):359-483, 2015.
[17] A. Gupta, B. Eysenbach, C. Finn, and S. Levine. Unsupervised meta-learning for reinforcement learning. arXiv preprint arXiv: 4806.046400 2018.
[18] 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.

---

[19] 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.
[20] Humplik, A. Galashov, L. Hasenclever, A. Ortega, Y. W. Teh, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: 1905.06424, 2019.
[21] M. Igl, K. Ciosek, Y. Li, S. Tschiatschek, 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-
2019.
[22] A. Jabri, K. Hsu, B. Eysenbach, A. Gupta, S. Levine, and C. Finn. Unsupervised curricula for visual meta-reinforrcemet learning. In Neural Information Processing Systems (NeurIPS), volume 32, 2019.
[23] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations (ICLR), 2017.
[24] Kaddour, S. Samundsson, and M. Deiscnroth. Probabilistic active meta-learning In Neural Information Processing Systems (NeurIPS), pages 33:20811-20822,, 2020.
[25] L. P. Kaelbling, M. L.Littman, and A. R. Cassandra. Acting optimally in partially observable stochastic domains. In National Conference on Artificial Intelligence, 1994.
[26] L. P. Kaelbling, M. L.Littman, and A. C.Casandra.. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101:99-134, 1998.
[27] K. Khetarpal, M. Riemer, I. Rish, and D. Precup. Towards continual reinforcement learning: review and perspectives. In Journal of Artificial Intelligence Research, pages 1401-1476, 2022.
[28] D. P. Kingma and M. Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014.
[29] 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.
[30] L. Kirsch, S. van Steenkiste, and Schmidhuber. Improving generalization in meta reinforce- ment learning using learned objectives. In International Conference Learning Representa- tions (ICLR), 2020.
[31] 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: 19:99884-19895 2020.
[32] 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.
[33] K. Lee, K. Lee, J. Shin, and H. Lee. Network randomization: A simple technique for generaliza- tion in deep reinforcement learning. In International Conference on Learning Representations (ICLR), 2020.
[34] K. Lee, Y. Seo, S. Lee, H. Lee, and Shin. Context-aware dynamies model for generalization in model-based reinforcement learning. In International Conference on Machine Learning (ICML), pages 19:5757-5766. 2020.
[35] 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:27222- 27235, 2021.
[36] Z. Lin, G. Thomas, G. Yang, and T. Ma. Model-based adversarial metareainmmeeennnnrremmmmm learning. In Neural Information Processing Systems (NeurIPS), pages 33:10161-10173, 2020.
[37] S. Liu, A. Davison, and F. Johns. Self-supervised generalisation with meta auxiliary learning. In Neural Information Processing Systems (NeurlPS), pages 32:1679-1689, 2019.
12

---

[38] J. J. Martin. Bayesian Decision Problems and Markov Chains. Wiley, 1967.
[39] B. Mehta, T. Deleu, S. C. Raparthy, C. J. Pal, and L. Paull. Curriculum in gradient-based meta-reinforcemenn learning. arXiv preprini 200.v:200..07956 2020.
[40] L. C. Melo. Transformers are meta-reinnforcement learners. In International Conference 2022.
Machine Learning (ICML). pages
[41] R. Mendonca, X. Geng, C. Finn, and S. Levine. Meta-reinforcement learning robust to distribu- tional shift via model identification and experience relabeling, arXiv preprint arXiv: 0006.07178. 2020.
[42] T. Nam, S.-H. Sun, Pertsch, S. J. Hwang. and J. J. Lim. Skill-based meta-reinforcement learning. In International Conference on Learning Representations (ICLR), 2022.
[43] R. Raileanu, M. Goldstein, Szlam, and Fergus. Fast adaptation to new environments via policy-dynamics value functions. In International Conference on Machine Learning (ICML),
pages 119:7920-79935 79371 2020.
[44] R. Raileanu, M. Goldstein, D. Yarats, I. Kostrikov, and R. Fergus. Automatic data augmenta- tion for generalization in reinforcement learning. In Neural Information Processing Systems (NeurIPS), pages 34:5402-5415, 2021.
[45] 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-6566 2017.
[46] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off-policy meta-reinforcemmnnt learning via probabilistic context variables. In International Conference on Machine Learning (ICML), pages 11.5331--3400 2019.
[47] D. Rao, F. Sadeghi, L. Hasenclever, M. Wulfmeier, M. Zambelli, G. Vezzani, D. Tirumala, Aytar, J. Merel, N. Heess, and R. Hadsell. Learning transferable motor skills with hierarchical latent mixture policies. In International Conference Learning Representations (ICLR), 2022.
[48] Rothfuss, D. Lee, Clavera, T. Asfour, and P. Abbeel. Promp:
Proximal meta-policy search.
In International Conference on Learning Representations (ICLR), 2019.
[49] Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
[50] 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.
[51] L. X. Shi, J. Lim, and Lee. Skill-based model-based reinforcement learning. In Conference on Robot Learning, (CoRL), 2022.
[52] X. Song, Y. Jiang, S. Tu, Du, and B. Neyshabur. Observational overfitting in reinforcement learning. In International Conference on Learning Representations (ICLR), 2020.
[53] Y. Song, A. Mavalankar, W. Sun. and S. Gao. Provably efficient model-based policy adaptation. In International Conference on Machine Learning (ICML), pages 11:9088-9098, 2020.
[54] B. C. Stadie, G. Yang, R. Houthooft, X. Chen, Y. Duan, Y. Wu, P. Abbeel, and I. Sutskever. The importance sampling in meta-reinnforcement learning. In Neural Information Processing Systems (NeurIPS), pages 31:9280-9290, 2018.
[55] L. Sun, H. Zhang, W. Xu, and M. Tomizuka. Paco: Purameter-oomposstionall multi-task reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2022.
[56] Y. W. Teh, M. Jordan, M. Beal, and M. Blei. Sharing clusters among related groups: Hicrarchical dirichlet processes. In Neural Information Processing Systems (NeurIPS), 2004.
13

---

[57] J. Tobin, R. Fong, A. Ray, J. Schneider, W. Zaremba, and P. Abbeel. Domain randomization for transferring deep neural networks from simulation the real world. In International Conference on Intelligent Robots and Systems (IROS), pages 23-30, 2017.
[58] E. Todorov, T. Erez, and Y. Tassa. Mujoco: physics engine for model-based control. In International Conference on Intelligent Robots and Systems (IROS), pages
2012.
[59] 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.
[60] J. X. Wang, Z. Kurth- NeltonNelson 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.
[61] K. Wang. B. Kang. J. Shao, and J. Feng, Improving generalization in reinforcement learning with mixture regularization. In Neural Information Processing Systems (NeurIPS), pages 0177968-7975 2020.
[62] 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.
[63] R. Wang, J. Lehman, A. Rawal, J. Zhi, Yi, J. Clune, and K. Stanley. Enhanced poet: Open- ended reinforcement learning through unbounded invention of learning challenges and their solutions. In International Conference on Machine Learning (ICML), pages 11:99400-999 1951,, 2020.
[64] S. Whiteson, B. Tanner, M. E. Taylor, and P. Stone. Protecting against evaluation overfitting in empirical reinforcement learning. In JEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL), pages 120-127, 2011.
[65] 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.
[66] 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.
[67] M. Yin, G. Tucker, M. Zhou, S. Levine, and C. Finn. Meta-learning without memorization, In International Conference on Learning Representations (ICLR), 2020.
[68] M. Yoo, S. Cho, and H. Woo. Skills regularized task decomposition for multi-task offline reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2022.
[69] 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.
[70] 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.
[71] T. Yu, D. Quillen, 7. 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.
[72] C. Zhang, O. Vinyals, R. Munos, and S. Bengio. A study overfitting in deep reinforcement learning. arXiv preprint arXiv: 1804. 06893, 2018
[73] L. M. Zintgraf, K. Shiarlis, V. Kurin, K. Hofmann, and S. Whiteson. Fast context adaptation via meta-lcarning. In International Conference on Machine Learning (ICML), pages 97:7693-7702 2019.
[74] I. M. Zintgraf, K. Shiarlis, M. Igl. S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. Varibad: very good method for bayes-adaptive deep RL. via meta-learning. International Conference on Learning Representations (ICLR), 2020.
14

---

A Pseudocode
Algorithm Subtask Decomposition and Virtual Training (SDVT)
Initialize Encoder 9r- decoder ps. policy Ty. set training tasks Muxi, virtual ratio pv. GMVAE buffer BVAR, policy buffer Bpo, total number of meta-episodes to train on Pameas number rollout episodes per meta-cpisode Rooll running mean of subtask composition y.
for meta-episode & O, ...., 7lmata do Sample training task Mr pmMimin)
Sample virtual meta-episode flag Bernouli(pp.)
V = then
Sample an imaginary subtask composition ÿ Dirichlet(W)
end if
Reset Au. yo, Bpot for timestep = ...... Aholl x H
if mod H = O then
Reset rollout episode end if
To)
Sample an action a, m-s(-131 Mie (les Me), aoe (be, v1), wa, (hc))
Take an environment step
8++1 ~ Tr((-(s,, 0z)
rt+1 4- (u 0e, 8c+1)
if V then
Σs NNMe, (he, 1)) 03. (he, ())
R. Pog (Ex)
Pet1 pe (-D(--DO(ss [st, as, sc+1] Re)
Replace T++1 f- f++1
else
Add the transition (8t, at, 8641, F(41) to BYAR ... VAE NOT TRAINED WITH VIRTUAL DYNAMICS
end if
Add transition (sr, dr, (141) to Bpal
Update hidden embedding he+1 = 9s, (T:+1)
Update subtask composition y Cat (way (be41))
Update the running mean subtask composition y RumningMeanUphate(( M1+1)
end for
Update GMVAE (0, 0) {s. 0) V (0.9) ΣE ELBO, with samples from BYAE Update policy z6 &- vi + & with samples from Bxa
Anneal virtual ratio Pa Po + Ap.
end for
B ELBO Derivation
The GMVAE's ELBO objective is derived as follows.
po (T Mc. 2t)
(d( [log pe (T. = Bd(M
log E (o
as (st, ze hel
po (T Uc, 2x)
E (m (h))
log (Mi, 20
log po (T 194, z (0)0 ( The )as (ze 2elhe, 9t)
(o(Mee )
E 90 (v
log.p((T. (2x) + log po qs (ee W.)
p(ve)
E (a
+ log
(o hz)
(8)
Eq. ( equivalent to the ELBO objective Eq. 3 without weighting coefficients. We assume
that the reconstruction T.H+ is conditionally independent of the subtask composition yr given 2t-
15

---

Implementation Details
C.1 Reference Implementations
MAML, RL2, and PEARL To replicate the results of MAML 13 RL 9. and PEARL 44 as
reported the Meta-World paper [7I], we utilize the exact version 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 leak-inffrence-baasee methods (SDVT, SD, LDM [35 and VariBAD benchmark based on the VariBAD's implementation We
(74) are adapted to the begin with VariBAD's hyperparameter configuration for MuloCo (88 Ant-goal and modify certain hyperparameters to accommodate the Meta- World benchmark (e.g.. batch size, network capacity, etc.). In line with VariBAD and L.M,, we train SDVT's policy using PPO 49). The GMVAE is based on an implementation that employs the Gumbel-Softtax 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 comprehensive list of SDVT hyperparameters. shared among SD, LDM, and VariBAD ensure fair comparison, please refer to Appendix C.3 com/ suyoung /SVT
and our source code available at
SDVT GMVAE
LDM VAE (recurment wncnder
Gaussian mixture
(recurent encoder dynamics decodar)
dynamics deceder!
Policy network
Policy network
Virtual training (subtask)
Virtual training jparamocricl
Virtual training (subiask)
Virtual training (parametric]
SD GMVAE freourrent encoder dynamics decoder)
VariBAD VAE
RL2
Gaussian mixture
Dynamics decoder
(recureent encoder dynamics decoder)
Recurrent encoder Policy network
Policy network
Policy
Figure 5: Schematic overview of algorithms. SDVT without a Gaussian mixture reduces to LDM, LDM without virtual training reduces to VariBAD, and VariBAD without VAE decoder reduces to RL2. MAML and PEARL use fully connected networks.
C.2 Computational Complexity
Table 2: Computational complexity. The total wall-clock time required to generate the results for the ML-10, averaged across all eight random seeds.
SDVT SDVT-LW
SD
SD-LW
RL
MAML
PEARL
VariBAD
LDM
Wall-clock time (hours)
142
140
138
135
192
17
258
126
131
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. Table 2 we detail the total wall-clock time expended by 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
http:://ggtaabb.cc/mmmmiigggggga httpss//ggthubb.co///rieeeeee/
16

---

evaluations, and the time taken to train the neural networks. Despite incorporating the GMVAE and virtual training, our method's computational demand does not substantially surpass that of VariBAD.
C.3 Hyperparameters
Table 3: Hyperparameters of SDVT and SD. Hyperparameteers of SDVT used for Meta-Words MI-10 and MI.-45 along with the notations in the manuscript and the argument names in the source code. SD shares the same hyperparameters except without those of virtual training. Different hyperparameters of our lightweight variant are denoted (-LW).
Category
Description
Notation
Value (ML-10, 45) Argument Name
Rollout episode horizon
H
500
max_epi
nax_episode_steee
Number of rollout episodes
Thril
10
max_rollouts_pex, _aa_ttaa_pppaakk
General
Discount factor
0.99
policy_gamm
Number of parallel processes
10
sum_processes num num_frames
Total Environment Steps
2.5e8, 4.0e8
Optimizer
Adam
optimiser_
Learning rate
1r_vae
Subsample ELBO
100
vae_subsample_elae
Subsample decodes
100
rae_subeamplee_eccod
Buffer size (meta-cpisodes)
[BvEl
1000
size_vac ve__vaeeeuuferrr
ELBO categorical coefficient
Oe
1.0
cat_loss_coeff
BO categorical coefficient (-LW)
Oe
0.5
cat_loss_coeff
ELBO Gaussian coefficient
0gg
1.0
ELBO reward coefficient
10
rew_loss_coeff
GMVAE
ELBO transition coefficient
OT
1000
state caste_loas__oeeee coeff
ELBO occupancy coefficient
EL
Oa
occ_loss_coeff occ occ_loss_coeff gumbel_tempereeee num_ _ad_tee vae ine_mixture__ee
ELBO occupancy coefticient (-LW)
aaa
Gumbel softmax temperature
temperature
Number of updates per epoch
10,20 20
Subtask dimension
K Nirin
10,45
Subtask dimension (-LW)
K
dee_mixture_mnie vae
Gaussian dimension
dim(z)
5, 10
latent_
Dropout rate
Plop
0.7
dropout_ran
Reward reconstruction objective
MSE
rew_pred_type
State reconstruction objective
MSE
state_pred_type
Algorithm
PPO
policy
Optimiser
Adam
policy.optimiser
Learning rate
Te
lr_policy
Optimizer epsilon
10c
policy_eps
PPO update epochs
5
ppo_num_epochs
Steps per policy update
|B
5000
volicy_num_stee
Policy
Number of minibatches
10
ppo_num_ming
PPO clipping parameter GAE
0.1
ppo_clip_paree
0.9
policy_tar
Initial standard deviation
1.0
policy_init_std
Minimum standard deviation
0.5
policy.min_std
Maximum standard deviation
1.5
policy_max_std
Entropy coefficient
1e-3
policy_entropy_comt
Initial virtual ratio
0.0
virtual_ratio
Virtual training
Virtual ratio increment per step
Ap.
5e
8, 2.5e
virtual_ratio_ineet
ELBO dispersion coefficient
10
ext_loss_coeff
17

---

Table Hyperparameters VariBAD and LDM. VariBAD and LDM employ identical hyper- parameters for the general and policy categories of SDVT and SD, as shown in Table 3 The sole distinction lies in the structural variation resulting from the utilization of GMVAE and VAE. We choose the Gaussian dimension that yielded the most favorable outcomes among 5, 10, 15, and 20.
Category
Description
Notation
Value (ML-10, 45)
Argument Name
Optimizer
Adam
optimiser_ve Ir_vae vae vae_subsample_et __elbos
Learning rate
Subsample ELBO
100
Subsample decodes
100
vae_subsample_ eedcddes
Buffer size (meta-episodes)
[Bval
1000
size_vas _aae__ae__tt Leee_aae_buffer buffer
ELBO KI. coefficient
0.1
kl_weight
VAE
ELBO reward coefficient
OR
rew_loss_coef?
ELBO transition coeflicient
orr
1000
state_loss coeff
Number of updates per epoch
10, 20
coo__vaeuuddatee
Gaussian dimension
dim(z)
5, 10
latent_din dropout_rate dropout_rate
Dropout rate (VariBAD)
Proo
0.0
Dropout rate (LDM)
pirop
0.7
Reward reconstruction objective
MSE
rew rew_ new_pred_type
State reconstruction objective
MSE
stade_pred__type
C.4 Network Architecture
The network architectures comprising our method are described in Table Prior to being input into the encoder or decoder, all state, action, and reward inputs pass through embedding networks. The values of K = dym and dim(2) differ for the ML-10 and ML-45, as indicated in Table
Table Network architecture. Details of the network architecture composing the GMVAE. The numbers in the Layers column represent the dimension of the hidden layers and the output.
Network
Notation
Architecture
Layers
Activations (last layer)
State embedding
MLP
[32]
(Tanh)
Reward embedding
MLP
16
(Tanh)
Action embedding
MLP
[16
(Tanh)
Recurrent encoder
94h
GRU
[256]
(None)
Categorical parameter
Weo
MLP
[512, 512, K]
ReLU (Softmax)
Gaussian parameters
Ho. 0 g=
MLP
[512,512, dim(=)
ReLU (None, SoftPlus)
Gaussian regularization parameters
[40 06.
MLP
[dim(=)]
(None, SoftPlus)
Latent context dispersion
p6,
MLP
[256, 256, 256]
ReLU (None)
Reward decoder
Por
MLP
[64, 64, 32,
ReLU (None)
Transition decoder
per
MLP
664, 64, 32, 40)
ReLU (None)
Policy
MLP
[256, 256,
Tanh (Tanh)
C.5 Aggregation method for the Success Rate
The main results the Meta- Meta-World paper (Table and Figure 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 these maximum scores across tasks is reported as the aggregated success rate.
This unconventional aggregation method does not accurately represent the meta-RI. objective, which aims to train single agent capable of solving multiple tasks. Meta- World calculates success rates for various tasks distinet time points during training, even though the agent might specialize different tasks at various stages. As result, the agent may excel at specific tasks by chance. Consequently, evaluating the agent more frequently could yield higher maximum scores. report the conventional final performance, which better reflects the meta-RL RL objective.
18

---

D Detailed Experimental Results
D.1 Performance on Individual Tasks
Table 6: Meta-Wore World V2 ML-10 success rate. We report the final success rates (%) of baselines and our methods for training tasks and test tasks of the Meta- World ML-10 benchmark. All results are reported the mean success rate + 95% confidence interval of 8 seeds 250M steps.
SDVT
SDVT-LW
RL.2
PEARL
VariBAD
Reach
10.0:1:1.00
22.03..5
53.8:11.5
56.2.19.11
19.
50.0
70.0
28.5
28.5144.1
Push
18.5
25.2416.5
61.24:33
68.8=10.6
60.0=13.0
0.000.0
0.000.0
25.8+12.9
12.
Pick-place
47.3:19.6
39.0+:5.4
52.54:..6
$5.0::3.99
42.3+2100
0.0:0.00
0.0:0.00
24.8+:3.3
27.2
Dave-open
100.0:0.0
100.010.0
93819.1
79.5122.11
10.2::19
86.81:22.7
62.5::33.55
5. Drawer-close
1.0.0:000
1000010.0
100.010.00
100.0:10.00
98.812233
99.5.1111
86.8.15.6
100.0.10.0
100.0=0.00
Button-press
98.8--2.3
98.8+-1..
96.24:99
100.0:10.00
87.548..3
93.04....
46.0+11.2
98.5=0.9
99.540.99
Pog-insert-side
31.2:t:13.2
24.5+11022
50,0:119.9
36.2=11.5
45.0=:21.6
0.0:0.00
0.0:0.00
19.8+14.00
21.8
Window-open
100.0=0000
100.0:0..0
98.8422.3
100.0:10.0
100.0:10.00
100.0:10.0
19.2
100.0:10.0
100.0=0.0
Sweep
96.2::4.88
72.24:29.0
23.8:222.2
13.8
75.0=17.3
0.0+0.0
0.0:0.00
60.24:32.4
61.24:2299
10. Basketball
91.2:1:6.4
78.8::9..8
13.8:114.3
91.2
43.8:116.6
0.0::.00
0.0:00.0
37.21:26.0
40.5+:28.9
Train
77.21.30
62.1.14.1
77.0...9
75.5.1..5
67.4144.4
42.2.14.5
23.2.199
88.2..899
16.7.112.3
Drawer-open
65.0:1:19.9
30.5::12.9
48.842344
45.0+224.0
2.2219
15.8+19.3
1.511
1284128
21.81220
12. Door-close
7.5.1.9.0
81.2.19.00
33.8.125.4
18.8
8.2.1.33
3.2.6.00
1.21155
27.0.1218
30.2.1.8.22
Shelf-place
0.0:10.00
1.0:1.22
0.0:1..0
022.0.2
0.0:0000
0.0:10.0
14. Sweep-into
90.0-18.5
51.21:18.9
71,2:11499
55.0:121.6
64.51883
0.0:0.00
0.8:::10
30.5::22.0
46.51:21.8
15. Lever pull
1.2.12.3
3.2.2.77
0.0:10..0
12.5
0.5.1088
0.5
0.5.10.9
0.2.00.5
0.5.109
Test meam
32.8:13.9
33.41:5.0
30.8:7.77
26.2:18.7
15.1:12.7
3.91:3.7
0.8:20.5
14.1
19.816.0
Table 7: Meta-World V2 ML-10 return. We report the final return of baselines and methods the last rollout episode on the Meta-Worid ML-10 analogous to the success rate in Table 6 Results are reported as the mean + 95% confidence intervals of seeds at 250M steps.
Index. Task
SDVT
SDVT-LW
SD
SD-IW
RI.2
MAML
PEARL.
VarBAD
LDM
Reach
3565-1132
3681#115
37084-228
3806+
2291-1276
3899-160
3841=5554
3829-1133
3740-159
Push
3923::339
3493::469
3636:1-288
357918800
815
17::1
15:22
2609-1900
2470-+11200
Pick place
2257 239
2240
213111441
2156.1593
501 101
61.0
511
1395
1561
Done-open
4441-:37
3171:1855
104
3530.1831
1249:111
3561:1753
1064-125
405111780
1134
Drawer-elose
4813=271
4850::10 10
4850:1 13
4854-9
2193:1336
4804:1:55
4081
2314
4841:118
4837
Buttou-press
3392=
3437-1133
32574:265
1189-111
956::300
3246:1161
140
Pog-insert-side
1903-:4330
1919:1344
263441588
2225:1592
652
7+:0
7+1
1148=7661
1237:-825
Window-opea
4486=:39
4371-277
4349::115
4398:179
1230:1112
3271#7332
784+:300
4485:132
4458:142
Sweep
4235::2202
3617:1682
382341702
3557:1634
906::223
791:38
51::14
2807=1260
2789-11333
10. Busketball 3549.129
3759 1.77
3466.1511
372611331
116
711
6.1.2
1111
2035 1256
Train mean
3656-162
3454 137
3630-2241
3525::297
11594-83
1822 136
1081
3055:1466
2963-1626
11. Drawer open
2558
2280.1221
1398
2310.1280
1649.1100
1737.1270
1130.187
2336.1.30
2389 1386
12. Door-close
566:1:335
807
1068-740
638:1670
$68
263--278
321
1121:1689
1343 1214
13. Shelf-place
501-:77
130
494=:102
434
578:170
0::0
0::0
211=114
328:+:044
14. Sweep-into
2221-1764
1455
1359-1632
1532-1721
490-1-46
41::3
48::33
658-3588
1 8899
15. Lever-pall
279-1:38
318:145
335-1-55
300:-488
291:166
156:b49
200-158
271:1:26
260-1-88
Test mean
1225:1160
1527-1214
1112:110
1043:1234
715.133
439-178
340:154
919-1 143
1166:12264
19

---

Table 8: Meta-World V2 ML-45 success rate. We report the final success rates (%) of baselines and our methods for training tasks and test tasks of the Meta-World ML-45 benchmark. All results are reported as the mean success rate 95% confidence interval of seeds at 400M steps.
Index, Task
SDVT
SDVT-I
SD
SD-LW
RI."
MAMT.
PEARL.
VanBAD
LDDM
1. Assembly
0.0::.00
0.0.0.00
0.0:10.0
0.0:1..0
2.2.10.4
0.0:10.0
0.0.10.11
0.0:0.00
0.0:10.0
Baskethall
3.0
29.8119.2
24.24:3..
21.019.8
0.0:0..0
0.0:10.00
19.8
10.24199
Button-press
95.0.1..6
94.5..0.6
100.0:10.0
9.701
9.
78.2.121.2
99.51.0.6
18.7
4ogdown
Button-press
94.0+7.
96540.6
99.0+
99.5+0066
98.5
78.24:21.6
10.6
99.8::005
86.2415.6
Hopdowa-wall
Button-press
93.2:43.3
96.5:20.6
89.018.11
99.0:10.7
84.2
95.51199
26.8111.3
98.011100
93.5
Button press
86.0=11.6
68.0:24.0
80.5
15.3
67.0=
11.2
20.24:3
74.5+:8.3
16.0=11.3
-wall
Collee-bun bution
100.0110.0
100.0.10.00
100.0:10.0
97.5
1.
92.014499
281113
78.0.114.77
100.010.0
Coffee-pull
37.2
33.24
45.5
46.5:1:2.00
4..0
0.5:10.3
0.0:1..11
$2.541
48.8 12.9
Coffce-push
64.8+66.5
15.0
56.822777
12.6
60.8=9.0
10.6
3.0+22.4
10.3
81.8
10. Dial-tur
87.2:1:4.4
86.2-22999
82.847999
67.2++11.9
14.5
81248.8
9.8:7.77
14.0
60.8:599
Disusemble
59.5=132.0
14.91:10.9
44.5112888
12. Door-close
100.0:1-0.0
100.0.1.0.0
99.5:10.6
97.0.13.6
98.2
1.0
52.2112.22
100.010.0
1.02
13. Door-open
56,2-31.2
56.0224.3
99.8-0.5
64.99226.9
54.0:223.88
6.8-119
0.5--0.3
98.0:1:
96.
14. Drawer close
100.0.10.0
99.0:11.2
97.2.1222
100.0.0.0
98.51.0.5
100.00.11
96.5.12.
99.2.10.7
96.
15. Drawer-open
96.5-11111
83.82
96.022488
87.2
83.5-844
0.0:0000
2.5:133
100.0:10.0
9322466
16. Faueet-open
99.5:t00.6
82,5+210
99.2=110
100.0=0000
868==9.7
10.5
18.8
95.5
18
17. Faucet-close
97.0:1:.55
17.0
99.2:10.77
97.0
56.8116.88
$2.8::22.5
15.0:6.33
99.5:10.66
99.2:110
Hammer
1.1188
6.8.5.88
0.010.00
19. Handle-press
-side
97.814.11
100.0:10.0
100.0:10.0
100.0:1000
99.820.11
83.2:11222
1281118
100.010.00
99.805
20. Handle-press
98.5
99.01112
100.0::0.0
100.0:1000
99.8
72.0+:20.2
66.8:19.88
100.0::.00
100.0:0.00
21. Handle-pull
40.0=22..3
5.2+5.2
50.5=25.2
39.5=24.1
83.2=4.3
13.8
0.8+0.77
41.24:23.3
31.8=28.5
-iide
22. Handle-pull
0.0:20.0
1.011.2
35.213155
3.8.2:.6
0.2:1:0.2
48.0
0.8:1044
1.0:0.77
1001100
23. Lever-pall
10.0.1..8
16.0:14.4
23.2.1:
3.2.1.99
20.0.111.3
24. Pog-insert
10::0.7
0.0:0.00
0.8:1077
0.5:0..6
1..5
0.0:10.0
0.0.1..0
0.0:10.0
2.2:12.7
-side
25. Pick-place
10.0=16.7
29.5
48.5+3.7
41.24
42.5=4.0
6.5+600
0.0:00.0
35.5::4.6
31.2415.9
-wall
26. Pick-out
of hole
16.4
23.2
53.0:17.0
43.8:8844
5.5::2.7
0.0:10.0
0.0::0.00
32.0::20.8
34.01:114
27. Reach
31.5416.7
13.5:4.
25.5
45.8-18.4
18.6
13557.22
27.0:14.7
28. Push-back
4088=21.3
31.0=1199
76.0=5.33
44.5
77.8=4.6
0,0::0.0
0.0+00.0
41.5+22.11
49.5=
29. Push
50.5++5.7
62.8=5.1
36.2-14.3
29.2+-18.
18.5-116
16.8:114.9
18::144
45.0:19.5
35.2=11.7
30. Pick-place
28.2:113.5
26.2:16.5
57.5
36.8
43.51:88
4.8:1..3
0.0:10.0
34.1:39
4..0.1.22
31. Plate-stide
75.8:15.7
56.5#6699
58.2+10.88
60.2+
20.3::14.4
17.0:113.2
9,5
55.2466.8
$7.0
Plate-slide
81.8.10.8
1.0.0.7
97.0
87.0.3.22
70.81...1
2.0:1..6
0.2.10.3
19.5
29.5
side
33. Plate-slide -back
82.0:662
86.0+99.7
83,524.2
76.5
90.0-3.9
0.0:00.00
0.8:00.8
84.0:66.1
80.2:::99
34. Plate-slide back side
71.8455.6
59.0=4455
60.8:11.2
73.5 10.8
81.5=119
5.5+:5.7
0.5+00.6
73.5+22
80.84:8.0
35. Pog-unplug side
10.2
10.81:2.0
54.0:1:.1
47.8
54.01:377
13.0:334
3,5:22.1
60.51::.88
48.8:19..1
36. Soccer
29.8
25.5:18.88
16.04:2.7
24.5
3.78
14.5
2.5:188
20.2455.4
15.24
37. Stick-pash
0.0=0.00
47.5-32.9
0.0:00.0
20.5::24.6
88.2-3.4
0.0+:0.11
0.0:10.00
0.0+0..0
0.0:10.0
38. Stick-pull
0.0=0.00
15.0=10.4
0.0:00.0
12.0:t14.4
4.54::.6
0.0:::000
0.0+:0.0
0.0:00.0
0.0=0.00
39. Push-wall
51.011:99
50.0
62.8-121.77
27.51:20.8
72.5
18.2.1..5
0.2:102
63.8
24.0::20.5
40. wall
16.2.1.59
7.011.99
27.0.18.1
28.0.1..3
6..0
35.221.0.22
15.0:6.9
18.0.13666
13.8.19.1
41. Shelf-place
0.0::0.0
0,5:0066
0.8:100.7
0.0:10...
3.8:1:99
0.0:10.0
0.0:10.0
4.5:14.6
0.:11:00
42. Sweep-into
8.55
97.8.12.66
43. Sweep
11.8=10.11
27.0=174
33.0
20.0:::3.2
41.54-1111
0.0:1:.00
0.0:10.00
40.54228.1
5.5-6.55
44. Window-open
99.5+:0.6
100.0::0.0
100.0=0.0
100.0-0.00
97.8-100
91.5#66.6
27.5=11144
98.2411.6
98.8+22.3
45. Window-close
100.0:10.0
99.0=11.2
100.0:1000
100.0:0000
95.2=333
97.0
27.2+11.0
100.0::.00
99.8:0055
Train
55.6.14.2
50.4:1.11
5..7.1..5
58.010.4
32.011.4
1033124
$7.0.1.22
$4.1.009
46. Bin-picking
U.0::0.0
1.5188
22211111
1.2:110
124:0.9
0.0:10.0
0.0::0.0
18:1111
0.5:0099
47. Bax close
0.5:10.6
0.0:10.0
0.510..6
0.0.1..0
0.510.3
0.0.10.0
0.0.1..0
35114
0.0:0.00
48. Hand-insert
2.2-2.3
3.823.11
L02112
15+18
5.24:3.0
23.0+:15.8
0.0::000
1.51111
08:100
59.2:199.1
74.8:17.5
49.8115.55
51.519.5
11.2:110.5
37.5.10.33
50. Door-unlock 78.5:1.7.9
75.8:19.1
61.8:10..2
72.5
38.0--11.3
64.5.119.2
25.0:19..2
642.4..3
63.2.699
Test
28.1:13.2
31.2:11.2
23.0:1:.11
25.4.2.99
11813:22
19.8:6.33
6.71133.
22.11:3.5
24.8:199
20

---

Table 9: Meta- Meta-World V2 ML.45 return. We report the final return of baselines and our methods at the last rollout episode on the Meta-Worid ML-45 analogous to the success rate in Table g Results are reported as the mean + 95% confidence intervals of seeds at 400M steps.
Index. Tisk
SDVT-LW
SU-LW
RL
MAML
PEARL
VariBAD
LDM
1. Assembly
481314885
2573
2314
2285.1195
1202.19 19
268.170 100:90
208.19
1434.1161
2572.1.6
Baskethall
688
14434:224
22151148
129
4891:38
6:2
1311:681
2227:236
Button-press
3120=t447
3297=328
3580=87
1343+1176
1344
2660=633
543#273
3348:2227
3208+3544
Hoodown
4. Button-pre press
4opdown-wall
3217:3384
3386::24
3144+110
1892+185
1359+18
26642238
604::144
3337=2279
3218:3319
5. Button-press
3100=:270
2977 199
2833-2208
2038 193
1165=29
2942 178
434 195
7878
2745+2205
Button-press
3401:1:220
1916:194
3148:157
3165-992
1378:1:7
2589
327 183
3040 136
10R
-wall
Coffee-button
1605=6678
2196::433
2517::296
2420=+:003
1348::40
3342=157
80:220
2253++303
2610::271
Colfee-pull
1032:4:007
9161:323
1010:1338
1305::395
572::355
38:77
1292::260
1244-1-220
9. Coffee push
1424:1.44
889 4455
1254
1687.1497
504
111
20:110
1901
491
Dial-turn
1821::478
1704--520
18354-414
155342544
44914007
130
1101--571
11. Disassemble
2451=12226
8974350
3387+3882
4840==4885
1642-13977
130=58
150+
1564+749
2268+688
12. Door-close
4178:::204
4404:222
43
4251
1231+220
4265=1188
2430+9965
4181=1105
4067:1174
13. Door-open
3086:761
22154:430
4134+151
2833#517
13124216
1095=316
149-112
373
3857+3325
14. Drawer-close
4310:12274
4544=54
4392:::007
4680-880
1626:
4725:128
4539:
4515=1112
4506:1:206
15. Drawer-open
39641:96
36241:334
42651:106
4015:1:221
2259
4180:1114
3876.1270
16. Faucet-open
4572-168
4057-1692
4582
4674
1925 193
3757
1810:1497
153
4479.1118
17. Faucet-close
4166:2351
3656+471
1457
4300:195
1662+1115
3298-561
1470:-414
4267457
190
18. Hammer
466::60
46
469
159#16
738+-233
149
420
476
184
19. Handle-press
1976#192
4722#11
4767=
4797=17
2276
3696#522
652=524
19
4737#:51
-ide
20. Handle-press
4601::74
183
1480-1147
4789::14
26
2902:1791
3036:19588
4728
4747:1:55
21. Handle--ull
-ide
853-1523
83.130
1507
907.1331
5651143
199.1.00
12:15
991:1576
827
22. Handle-pull
153
1491=110
2260::666
1541=100
2048+86
2001
29:119
1549=557
1612:1170
23. Lever-pull
446:884
496:774
4834:39
500:1106
305+:1
264260
119::38
457#:32
345::227
24. Peg insert
767-2408
757#:48
1106+1186
1967#192
640
33#:24
7=2
178=101
962#2232
-side
25. Pick-place
1647.1893
20111:717
24
2570.1543
491
3899
0.10
2042.1501
1873
26. Pick-out
809::435
1015#365
1319+188
1385::310
270:
20
1
1129#
1247::392
-of-hole
27. Reach
2682::225
24481:13
3038-12388
3116:::82
2273:1:99
244
2466:11079
144
188
28. Push back
674.1362
514.1154
1456
1218.13334
433.142
11177
8:133
87111372
1404
29.
3130:1:290
31391665
2961+:253
2581+:530
9454:32
6341:555
34:13
185
3524::200
30. Pick place
Pick-place
1221:1650
1582.1433
159
122
$76.123
247.1.222
511.1
2034
2128.1144
Plate-slide
3526:1161
2950-1212
3022:::96
2838:1303
2863-11349
1043
561
2858
2801
32. Plate-slide
3058::289
13622:7
3471::258
2761#3117
949::20
3224:249
56::35
3107=1167
2351+:462
side
33. slite
3783+2204
3790+:239
38444:128
3754+2272
1544+10
581+:644
655-2259
3900-1155
3720=1325
hack
34. Plate-slide back side
3732
376411149
3755
3987.1228
1561.17
67611420
133.156
4099
1094
35. Peg-unplug -ide
082-2390
971
11871111
1079::464
401::36
62:4:6
45:128
1523:+142
1087+2282
36. Soceer
1374:1169
1566:1:007
1311::129
1348::332
671:131
439::359
98:160
1399-190
148611-2
37. Stick push
47.170
78631533
11.12
291.1333
955.14
1281122
6.1
17.12
10.14
38. Saick-pall
25-1-33
908-1616
144:3
600-702
73011-46
39:1:52
7
1911
11::4
39. Posh-wall
2547::426
2197=:088
3291:1-211
2496
877485
597+:532
20:17
2813:7116
2744-1:549
40. Reach-wall
1871-:157
1889-1412
2447-1-662
2741
2303-189
2802-1: 1082
1741
2102-:
183611569
41. Shelf-place
608-:308
657-:195
920-4555
1056:1165
623+661
94:8
0=:0
195+118
125
42. Sweep-into
3755:+333
3847=2400
3869-2214
3584:1495
7934:67
6234-541
148:182
3750::159
4151+-185
43. Sweep
1148:198
1433:1477
190314449
1415
872.1/300
160:182
36.19
18651179
121911325
44. Window-open
3789-1.249
4004 190
4138:156
4121
1163-137
2463-L 1004
184
4177
198
45. Window-elose
4201
4333 135
4293-b2009
4496::29
1152: 109
3215:4225
990:1417
4346
4450:1:33
Train
2379:t214
2294+:202
2672::79
2578-264
1411=:22
1388=1104
597=1121
2492
2515::67
46. Bin-picking
114:1:44
113:::56
322:1:70
140::72
265:1:33
37:16
12:15
127::27
90:1:99
47. Box close
190.1.4
128.12
143.15
137.99
1
115.1:6
176.1.0
243.183
151110
48. Hand-insert
299-164
465-1:01
386:125
392-164
467-1-88
387-4-88
3774-1
398::39
49. Door-lock
1694+:302
1973=:160
1359+2262
1427+:289
995-23033
8214:243
960=1499
1358:1445
1573+1168
50. Door-unlock
1898:1158
1792-1:225
1721-:132
1870:E:245
769-1108
1880:1-426
1374:::537
1707-4324
1627:1:199
Test meam
839.1.74
194
786.1.69
793:1-49
100
658.196
506:1
762.1.40
768.163
21

---

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 MI.-45) are used to present the main results in Table D
100
SDVT
SD
RL2
PEARL
LDM
SDVT-LW
SD-LW
MAML
VariBAD
Success Rate (%)
60
40
20
0.0
0.5
1.0
1.5
2.0
2.5
Steps
Figure 6: Learning curve on ML-10 success rate.
5000
SDVT
SD
RL2 MAML
PEARL VariBAD
LDM
SDVT-LW
SD-LW
4000
3000
Return
2000
1000
0.0
0.5
1.5
2.0
Steps
1:8
Figure 7: Learning curve MI.-10 10 return.
22

---

100
SDVT
SD
RL2
PEARL
LDM
SDVT-LW
SD-LW
MAML
VariBAD
80
Success Rate (%)
20
0.0
0.5
1.0
1.5
2.0 Steps
2.5
3.0
3.5
4.0
Figure Learning curve on ML-45 success rate.
5000
SDVT
SD SD-LW
RL2 MAML
PEARL VariBAD
LDM
SDVT-LW
4000
3000
Return
2000
1000
0
0.0
0.5
1.0
1.5
2.0 Steps
2.5
3.0
3.5
4.0
Figure 9: Learning curve on MI.45 return.
23

---

D.3 Subtask Compositions Learned over Training
0:22
6:20
6.30
3:00
0.68
0.28
6:33
820
220
326
6.25
1.21
3.0
0.2
6.19
5:30
0:16
03
0:33
6:9
31
0:33
0:33
6:23
833
330
3:3
0:19
2:13
333
0.33
6:28
5:07
2.1
0:27
00
3.00
0.06
RI0
6.00
8.00
0:35
021
02:
6.1
5:11
1:1
0.3
6.2
6:21
621
021
521
6:21
5:21
921
9:21
0.00
6:37
0.09
651
$500
Seed aT SAM stree
0.
01
6:07
5:0
R::33
D:11
5:95
3P1
0:00
0:03
RLE
600
5.00
10:00
3.00
0:30
6:30
6.03
3.50
6:22
9:53
3.41
0
0.50
6:26
A.45
2.03
2:20
3.00
0.00
2324 stE
2
0-21
6:22
1::2
1.00
6:32
E.04
0.00
0.5
0.00
6.24
F:03
3,00
0
E-
6:00
E27
000
3:33
8:33
0:00
0.06
8:33
8.00
100
3.13
6:79
5.03
3.00
0.00
6:22
R/02
0:24
E
3.00
0LEE
E00
2:00
0:00
6:03
8:0
5:00
6:21
5.91
3.00
0.00
6:N
S
R
600
1.00
3.00
6.30
5.03
300
100
6:30
1.22
3.03
0:37
0:50
18
11
12
15
14
15
39
11
3>
(a) SDVT-LW (K = 5,00 0.5, 0.5,00 = 0.0).
Sand
0:31
8:39
1:0
6:18
F:15
530
8:37
0:07
2.39
10
6:32
6:10
5:10
3:10
0.16
1.31
6:33
8:10
310
3:10
6:15
5.15
010
012
6:27
5:07
3.04
0.16
0.24
6:27
6:1
111
3.12
0.08
E07
0.21
0:32
6.39
5.10
1.00
8.16
6.35
6.30
6:30
5:10
3:00
6.18
5.13
0100
0:10
010
6:38
5:08
3.07
10C
0:32
9:24
8:37
1:2
3.10
0:38
697
696
5.66
0.06
310
6.18
0.2
0.0
6.3
6.12
3.00
0.0
0:22
0:21
8:30
0:30
3:00
330
0:3
6:33
830
230
210
6.28
5:33
330
022
6:25
F.15
235
0.00
0
8:26
2:30
635
0.25
0
0:2
6:00
6:10
3:10
0.38
0:38
R:23
830
330
2:6
6.18
F:13
9:0
0:37
03
6:33
6.12
2.3
9
0.2
6:21
6111
1122
3322
0:33
611
831
0:23
0:2
5.10
0.15
A.15
6:33
610
510
6:29
5:15
110
0:12
0.00
5:19
512
are
0.22
1.
3.1
6:17
R17
6:11
515
6-12
0:12
3:10
6.55
2:14
38
11
12
13
14
35
39
11
12
34
35
21200
Seed
E.D0
0.01
0.30
0:11
0.0
0.33
10
100
000
1.69
1:22
2:11
0.68
0.51
6:33
E34
226
2.37
6.28
8.02
3007
0
0.6
6:23
6.35
2000
0.00
62L
F22
22
003
F33
0:52
0:2
0:23
6:33
9:58
0.08
0.04
0:25
8:0
206
3.55
8:72
2.03
3.04
0.05
6:70
5:02
2.00
100
0
0:27
3:34
100
300
000
E:00
8:00
0:00
0:22
612
9:32
6:31
617
2:2
9:0
6:29
F:1
912
0:34
0:00
6:2
R:1
9.03
0
8:34
F:1
835
0.e4
5:36
032
0.32
0.32
6:24
6.12
311
3:3
6.13
5.13
3.22
0:22
0.44
6.26
6.14
3,00
0.00
0
2:22
0.00
0.22
6.35
0.54
000
0:24
6:24
6.12
2000
0.32
0.24
0.21
6.4
214
3:3
6.12
2:14
214
0:24
0.00
6.39
6:22
2,00
0.00
0.04
0.23
6:20
2:21
3:28
0.42
611
6.10
5:39
0.27
611
120
314
6.17
F.15
735
0.33
0.00
6:31
5:28
9:3
9.00
9:39
0:23
Send aM
324M DA2S
Send Eat 2024 stook
30
6:53
3:50
6:32
9.03
3.00
0:00
0:30
1:00
9.13
0:53
7.00
3.00
8:00
0.04
6-31
300
6:20
6:22
3.0
2.36
6.26
5.03
3,00
0.66
0.00
6.28
5.05
3.02
0.03
6:33
611
200
100
1000
600
6:22
6:22
8:00
1:00
0.00
0.32
0.33
6.34
8:4
2.30
6.28
512
012
0.04
100
6.29
6.15
2.00
9.66
0:02
0:2
020
100
333
000
8:35
0.00
82
3.00
0.22
5.3
535
0.00
0:00
6.20
6.27
3.00
0.00
6:3
100
630
3:00
R:93
R3N
936
6:28
2.03
6:12
K111
8:30
0.66
R:53
6:14
3:58
6:27
6:11
0111
0111
0.0
A117
5:44
2/01
0
0:22
2:20
5:00
0:35
6.20
2.00
800
0.35
6.41
6.3G
2.00
6.18
0.24
0.00
6:20
6.13
3:30
0.00
[058
6:30
2:00
3.00
624
6177
4.64
1036
31 0swergee
12 14006 Ske
32 Leve
2.7
33 Cover
(b) SDVT (K- 10, 0. = 1.0, a. = 1.0).
Figure 10: Subtask compositions learned over training. We visualize the development of subtask compositions for (a) SDVT-LW and (b) SDVT during training on ML-10. Similar to Figure 4 each column represents the terminal subtask composition (yn=) averaged across 50 parametric variants. (a) As training progresses, subtask compositions of (4) "Door-open" and (12) "Door-close" merge due to shared transition dynamies, which may cause the "Door-close" policy to keep the door open during testing without virtual training, underscoring the crucial role that virtual training can play in these scenarios. (b) With occupancy regularization, as the training progresses, the decomposition process prioritizes occupying lower indices over higher ones in the subtask compositions.
24

---

D.4 Subtask Compositions of all Seeds
Soud2 +4 2222
252M
0:00
0:37
200
ZEO
6:32
8.02
227
0:20
0:42
EAL
2.03
3.00
2:5
3.00
24
L2
14
25
8:78
2.09
2/03
2.09
8:28
0:00
6.23
6:32
0.00
2:00
2.30
6.28
3.02
3.00
66
0.00
6:20
5.23
3
6.32
2:20
OF
0.00
6:30
2.00
3.58
6:38
5.93
3.00
0.66
633
622
SUR
0
6:52
2:0
3
IUSS
600
2014
0.66
3:06
6:12
0.02
1:00
5,05
3.03
0.66
0.50
5-30
5:00
3.00
002
6.00
6:00
0.66
8:33
3:00
0L0S
R:37
2.0
8:28
5:07
1:00
3.50
3.00
1.77
166
1.5
1
100
1:00
1.97
1.00
3.00
1.00
0:03
6:22
1.0
2.00
0.60
0:35
8:33
0.00
2:00
2.00
6:32
3.00
022
0.00
0:20
6:23
2.00
0.00
02
0.33
20
1:00
3.00
0.00
1:02
E.00
1.00
0.00
9:28
0:03
6-31
5:00
1:00
8-28
0:00
6:20
6/00
910
0.00
30
10
3/6
600
6.00
0:38
3:50
1.07
9:39
0:00
0:30
3
9
0:33
2:20
23R
T 25
0
0:07
6:54
D000
LCC
6:53
840
5:45
LSC
6:30
9:09
2.
D:44
0:24
5:28
5.00
0:05
0.53
3.00
KVD
8.00
10:15
0:00
6:93
0:0
1:50
F:LE
0-30
8:30
5:32
0:00
R:59
6.0
8::00
00
0.63
2.20
3:02
3.00
0.52
6:20
3.00
6:20
0:28
38
11
12
14
14
15
2
3
59
11
(a) SDVT-LW (K 5, 5,0 0.5, Go = 0.0).
0:06
0:50
6:30
0S3
100
3.00
6:33
5:03
3.00
0.04
0.50 6:30
5.33
2.14
0.60
0.
053
250
3.00
0.00
0,02
0.00
0.07
0:20
8:30
1.00
0.00
0,63
0:33
0.00
2:00
2.00
6:27
1.00
3.00
0:21 6:33
8:07
2.00
0.00
0.
0:23
2:20
2:00
3.00
0.08
002
0.00
832
0.00
8-31
8:18
6:21
3.00
0:66
6-34
6:2
8:00
2.00
0.06
0,02
0.37
8.0
R35
373
6:18
2.04
223
0.
0:00
6.20
8:22
220
0.00
0::
6.20
22
2:0
3.00
822
RGC
E00
8.00
0.03
0.27
1:20
0:33
8:31
835
2:00
3.00
6:22
2:23
0.50
6:23
8:35
3.02
0:22
6:37
5:20
1:16
0:32
632
[69
0:25
0
0:03
6:34
1:00
6:32
8:33
3:00
1:40
6:21
8:21
3.0
0:66
0:0
6:28
6/22
022
6:22
6:20
2:00
3.6
602
600
10
6:
2:30
0.45
6:33
3.21
6:26
2.05
[83]]
0.60
0.00
6:36
5.05
3.21
6.33
5:21
00:
6.00
0-42
38
Seed
x 25RM maas
Send kax 2524 steor
[-21
LED
8:10
E311
1100
8:10
2303
E101
210
110
1003
D10
0.06
6.33
E33
3.34
6.23
1:03
3.00
02
0:00
6.12
3.00
0:
6:24
8:50
0:00
0:07
[63
1:00
3.00
0.00
00
0.33
0.00
2:00
3.50
0:38
200
3.00
0.00
6:49
2.02
0.00
0
0:25
0:20
2:3
535
0.08
007
000
888
0:00
8:36
E.05
1:10
0:33
612
5:00
0.63
200
3:00
6:33
0.00
6:24
5:22
0/6
6:23
8-30
3.00
0.0
023
0:20
6:38
E24
330
E.24
2:03
[230
28
0.00
6.26
6.21
235
6:27
3:20
[62]
10:5
6:30
2:30
3:00
IA34
8:03
3.50
0.0
0.85
2:00
3.00
2.00
3304 aoaza
Seet
2322 stera
0.00
6.33
1.00
2.20
6.23
3.28
3.01
0.02
0.00
6.20
5.03
5.1
9.00
0:
0:53
0:30
1:10
066
1:22
1:31
8:29
(4:33
111
1.18
6:38
0.31
3107
0.26
0:200
6:18
5:34
3:23
6:21
8-31
3.60
6:12
8:0
3.20
6.38
5.03
3.07
0.02
0.00
0.27
6.02
3.03
0.5
0:
6:30
0.44
000
1:35
6:23
6:15
2:00
0.06
8:33
0:34
E62
3.0
6:25
2.2
0:05
6.26
1:23
2.04
2.00
02
6:25
8:35
3.04
0.08
E00
E27
0:26
3.5
0.3:
0.24
1113
0.00
0:34
6:34
642
3.0
6:28
5:02
33
9:32
040
0:13
8:25
215
0
6.22
5:36
1:
3:3
0.62
022
0.00
0:03
6:30
6-23
6:17
200
6-28
0.00
6:00
6:28
20:
0
0.00
8-30
3m
600
0.03
0.20
0.00
0.66
6:23
6:3
8:00
0:43
6.28
5:03
3:28
0.05
0:32
6:28
0.24
331
3.66
0
0.27
2:34
15
033
0000
0:12
14
L2
14
35
Swif4
Sead Far 2524
E:51
3:10
6-10
110
016
11:5
110
200
2:0
1:23
E-21
0:30
1:31 B10
310
11:0
6:18
5.11
3.00
0.60
6:200
6:18
5:1h
2/11
0.6
21
221
1027
(.00
3.06
1:00
6.13
1334
511
0.00
6:48
6.03
5.5
6X0
054
021
15-1
9:34
0:4
00
0:33
0:04
0:33
E03
200
6:25
231
235
0.
0.00
0.20
6.25
3X0
O21
F00
100
100
100
E00
0:00
8:38
2.05
8:00
6:00
0.02
602
6:24
8:20
3.00
0.00
8.00
0.3L
821
283
2:20
6:33
2,02
2,07
0.00
0.0
6.22
6:44
2.00
0:23
6:25
2:00
3.00
1007
2.64 6.00
0.2L
10
6:30
0.00
2:00
3.30
0.22
5.05
2.01
0.05
0.00
6:20
2.02
[830
3.06
0.02
6.34
2.00
3.00
0.90
0022
6:05
1o. sworaa
12 D.DDE Lssarage
12. Lanee 1.S0MR-T
Dornge
35
(b) SDVT (K- 10, a. = 1.0, 0o 1.0).
Figure 11: Subtask compositions of all seeds. We visualize the subtask compositions of (a) SDVT- LW and (b) ML-10 after 250M training steps for all seeds. Decomposition processes differ among random seeds due to varying initialization and sample tasks during meta- vneta-trainnng. One subtask can be interpreted as combination of multiple subtasks and vice versa, depending on the seed and the dimension K. We rarely have a collapse to single Gaussian as in (a) Seed 5, which lowers the average performance. (b) With occupancy regularization, we find that the 10th element of the composition is rarely occupied, whereas the first element is occupied by many tasks.
25

---

E
Ablation Results
In order to confirm the source of the gain obtained from the proposed method, we conduct extensive ablation study based on SDVT and SDVT-LW on ML-10.
E.1
Occupancy Regularization
Scod at 250M stops
Sead 250M steos
5.00 00 60 00000000000000 00 0.00 0100 0000000000000... 00.0000000000000 2.000.0....000 0.00
0.00 00.00.00000 0.00000000000000000000000000.000 000 0,000,00000 0 50,0.00 00000 00000000000000000000000000000000 0.00 0.00 0.07
2.00 0.000000000
uod. 0000 oo 00
0.000.00 0.0 2.000000..000000
2.00 0.20 0.00 00 20 do 0.000
0.00
2.500.00 0..0..0..oo do 0.00d. non. 0.00 0.00 b.00p.000000
0.00
0.00 oe 0.000000.000 00.0 3.. 50 .0 0.000.000000000 00..0000..
0.03
2.50 3.33 0.00i, ooi, 6.13
8X330.00 0.0.0000000000 8:31 0.50
5.00 30
00 20 9.00 0.00.. 00.0000000
0.00
0.00020 0.00 0.000,000 0...0000.000 0,010.00 3.00000000.000
0.00 00
0.0.000.00.000000000 0.000000..0000.000000000000
00.0000000000000000
u.D
0.500.00
8.34 2.00
8:24
0.60
0:300:0:377 0.00000000000000.. 000..003.0000
0.12'
0.028.531..00
0.39
8.19
2.000.00
0-00
8380.00
0.50
2.09 00500 0.000 00...00... 00 0:440.00 0.00 5.00 00 5.00 5.00
8.00 0..40033000
10.00.00000
3.500.58
0.58 1811.0 0.000000 270 17.00 0..0033 0m1
Pokglace
Doorqpr
Swoep
Baskatbll
11. arannarepp
13 S0..hppa
35 tonuroo
A.Nush
1.Paa
Pakgloss
Decespee
Dmenncce core
Sawep
Dadestaal
11. 3amnsr-pp
32. 1...ancce
13 Shiiipace
14 Swessit
(a) SDVT (ae 5.0).
(b) SDVT (a. 10.0).
Figure 12: Occupancy ablation. We report the learned subtask compositions of SDVT (K
10, a 1.0) for different occupancy coefficients analogous to Figure 4b
Figures 4b and 12 indicate that our occupancy regularization effectively constrains the use of higher indices intended. We fix the default choice of a 1.0 that we found to work well on ML-10. Our default SDVT with a = 1.0 scores a test success rate of 32.8%, which scored the best among a {0.0, 1.0, 5.0, 10.0} that scored 22.8%. 25.5%, and 21.0%, respectively. Also, our default SD with a = 1.0 scores a test success rate of 30.8%, which scored the best among a, E {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: Ablation results. We report the mean success rate (%) and the difference caused by the changes made from the default SDVT-LW (K = 5, a. = 0.5, a, = 0.0) on ML-10 in parenthesis.
SDVT-LW
without
K
Qc
Dropout
Dispersion
3
7
10
0.1
1.0
2.0
65.6
73.0
60.1
69.5
70.8
59.8
70.3
66.3
ML-10 Train
(43.5)
(+10.9)
(-2.0)
(47.4)
(+8.7)
(-2.3)
(48.2)
(+4.2)
16.5
21.5
25.0
22.0
21.1
20.5
16.1
17.9
ML-10 Test
(-16.9)
(-11.9)
(-8.4)
(-11.4)
(-12.3)
(-12.9)
(-17.3)
(-15.5)
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 of the LDM paper (35), virtual training without dropout does not exhibit 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 substantial for VariBAD as it is for SDVT-LW's virtual training.
26

---

E.3 Decomposition Distribution
Table 10 emmphsszes the significance hyperparameters such as subtask dimension K and categorical coefficient aę in determining subtask compositions. is crucial to carefully set these hyperparameters based on the number training tasks and their correlations. our lightweight (LW) approach, we have selected K = 5 and or 0.5, which effectively distributes subtasks across the MI.-10 tasks. When Og is set to a small value, task classification collapses into a few subtasks, as depicted in Figure Ac. Conversely, with large value of a.. the entropy of the composition maximized, and all tasks exhibit a uniform probability distribution, as illustrated in Figure [d] However, it is noteworthy that SDVT-LW outperforms VariBAD and LDM in 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 oceupancy regularization.
E,4 Number of parameters
Table 11: Number of parameters and success rates. We report the number of parameters used by our methods and baselines. We demonstrate that our gain is not from the increased capacity.
Number of Parameters
ML-10 Success Rate (%)
Methods
Encoder
Decoder
Policy
Sum
Train
Test
SDVT-LW
1,047,455
174,821
235,401
1,457,677
62.1
33.4
SD-LW
1,047,455
25,637
235,401
1,308,493
75.5
26.2
LDM
502,580
25,577
202,249
730,406
56.7
19.8
LDM (matched)
2,144.692
175.593
267,401
2,587,686
64.2
22.0
VariBAD
251,290
25,577
202,249
479,116
58.2
14.1
VariBAD (hidden x2)
894,362
25,577
663,305
1,5832244
62.4
12.0
VariBAD (matched)
1,072,346
175,593
267,401
1,515,340
67.6
17.2
Table TI indicates that SDVT-LW employs and times more parameters compared to VariBAD and LDM, respectively. The encoder of our method has more parameters than VariBAD due the added categorical layer. The decoder of SDVT-LW includes the parameters the dispersion layers. However, 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 x2) While selecting encoder (GRU) and policy hidden sizes, we experimented with 128, 256, and 512, discovering that 256 works best for task-inference methods (LDM, SD only, and SDVT). Notably, VariBAD (hidden x x2) 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 '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
Table 12:
Masking contexts. We present the success rates (%) on MI.100, achieved by ablating the contexts fed into the policy, to illustrate the effective utilization of learned contexts by the policy.
SDVT-LWW masked (Wo,)
SDVT-LW
SDVT-LW
SDVT-LW
masked (Mo.. .o.)
masked (Woy Poo Co.))
ML-10 Train
62.1
58.6
36.9
34.8
ML-10 Test
33.4
33.6
25.8
24.5
27

---

We employ VariBAD's 's stop gradient architecture to avoid the instability that may arise from the concurrent training of policy and ELBO objectives. This may lead a potential vulnerability where the policy neglects contexts. However, in Meta-Worke 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 (wiby) or ( 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 is more critical than masking the categorical parameter of y. However, this does not suggest that y is redundant, given that the continuous context 2 is trained to include y's information. Interestingly, the test success rate, with all contexts masked (i.e., solely relying on the current state st), achieves 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 GMVAE to prepare for tests with unseen compositions of subtasks.
E.6 Smaller number of rollouts
Table 13: Success rate on ML-10 for smaller rollouts. We report the success rates at three different rollouts, Rolll 1, 2, 10. Note that the results reported in manuscript are for noll 10.
Training Success Rate (%)
Test Success Rate (%)
Methods
Plyll
Rrall
2
Thall
10
Tolll
Tolll
Thoil
SDVT-LW
61.4
62.3
62.1
32.8
32.7
33.4
SD-LW
74.9
75.1
75.5
25.8
25.2
26.2
VariBAD
57.0
58.4
58.2
14.5
15.3
14.1
LDM
55.5
56.0
56.7
18.4
18.5
19.8
We Rgoll 10, following the original setup of the Meta-World benchmark [69 7
Notably, this
benchmark utilizes dense rewards, therefore the context converges rather quickly, even within the first rollout (Figure 2b). As result, the performance across rollouts does not exhibit substantial improvement. To provide further insight, we report the success rates for smaller Pell in Table [3
28


Reviews:
[{"review_id": "6I8GCHQACy", "reviewer": "Reviewer_juhC", "summary": "Meta-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. ", "questions": "Could 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?", "limitations": "See other sections.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This 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", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for recognizing the many strengths of our paper, and we have taken the opportunity to address the concerns raised regarding the perceived weaknesses in our work as detailed below.\n\n### **W1.** Examination of Why SDVT-LW Outperforms SDVT\n\nWhat our model aims for is the decomposition of tasks into a mixture of subtasks (i.e., a multi-hot composition), facilitating generalization for test tasks that share common elementary subtasks. Table 10 in Appendix E.2 highlights the significant influence of the number of subtask types, represented by $K$, on the test success rate of SDVT-LW. Specifically, when $K=10=N\\_{\\text{train}}$, the test success rate decreases to 21.1\\%. This decrease in success rate from the optimal $K=5$ can be attributed to the increased probability that each task is classified into a unique subtask (i.e., a one-hot composition). On the other hand, when $K$ is too small, our model fails to learn an adequately diverse spectrum of subtasks. Consequently, we need to identify an optimal value for $K$ within the range of $1<K<N\\_{\\text{train}}$. \n\nOur introduction of SDVT with occupancy regularization was aimed at achieving performance on par with SDVT-LW, even without prior knowledge of the optimal value of $K=5$. Even when beginning with a sufficiently large $K=N\\_{\\text{train}}$, SDVT is able to constrain the higher dimensions of subtask compositions, although not to exactly 5 dimensions, as shown in Figure 4(a). Since SDVT requires some training epochs to diminish the higher dimensions, it lags behind SDVT-LW, which uses 5-dimensional subtask compositions directly from the beginning. Exploring ways to outperform SDVT-LW without knowledge of the optimal $K$ or even without the number of training tasks, $N\\_{\\text{train}}$, presents an intriguing direction for future research. We will provide a clearer explanation in the main manuscript regarding the distinctions and the rationale for introducing both the SDVT and SDVT-LW models.\n\n### **W2.** Relevant Works\n\nWe are grateful to the reviewer for introducing relevant works that align with our research. These references will enhance the related works section of our manuscript. Exploring the integration of techniques developed within the offline skill-based setup as a part of our subtask-based methodology presents an interesting future development.\n\n### **W3.** Confusing Term \"non-parametric\"\n\nWe apologize for the confusion caused by the use of unclear terminology. We appreciate the reviewer for suggesting alternatives. Please refer to the global response **G1.** for more detail."}, {"review_id": "e95ZyFcNXN", "reviewer": "Reviewer_mp13", "summary": "This 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.", "questions": "- 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?", "limitations": "Limitations of the proposed method are not discussed at a sufficient level of detail.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "## 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.", "weaknesses": "## 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for recognizing the many strengths of our work. In response to the concerns raised, we outline below how we plan to improve the presentation of our work. We are excited to incorporate your valuable feedback, thereby contributing to the RL community more effectively.\n\n### **W1.** Clarifying the Contributions\n\nWe acknowledge that the introduction of multiple components in our work may have led to some ambiguity. To shed clearer light on our unique contributions, we will add a paragraph in the Method section, detailing the following aspects:\n\nForemost, our primary contribution is the proposal of SDVT-LW. This model leverages a Gaussian mixture encoder for contexts, effectively parametrizing tasks by their subtask compositions, and carries out virtual training within the realm of these compositions.\n\nFurthermore, we propose SDVT with the occupancy regularization strategy. This generalizes SDVT-LW to adapt to more difficult conditions where there is a lack of prior knowledge of the optimal number of subtasks $K^*$. We initialize with a large value $K=N\\_{\\text{train}}$ and employ occupancy regularization to downscale higher dimensions, navigating to discover the most efficient number of subtasks even without prior knowledge.\n\nFor the purposes of virtual training, we've adopted two methodologies that have found application in previous studies: dropout [2] and dispersion via structured VAE [3]. 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. \n\n### **W2.** and **Q1.** Occupancy Regularization\n\nPlease refer to our response to **W1.** and the global response **G2.**.\n\n### **W3.** Clarification on Ablations\n\nWe respectfully wish to direct the reviewer's attention to the fact that we have indeed provided ablations concerning occupancy regularization, dropout, and dispersion in Appendix E. In a brief summary, and consistent with the findings of prior research (Appendix E of LDM [2] and Appendix A of DiAMetR [3]), dropout and dispersion are indispensable elements when generating virtual tasks with unseen dynamics. Dropout and dispersion allow extrapolated generation over unseen (states, actions) and contexts, respectively. Occupancy regularization plays a crucial role in enhancing test success rates of SDVT, even when starting with a large subtask dimension. To improve accessibility and comprehension, we will incorporate a summary of the function of each component into the main manuscript, as also outlined in our response to **W1.**\n\n### **W4.** Weak Analysis\n\n- Subtask Compositions: We have showcased all tasks as matrices to highlight the differences in distribution for various hyperparameters at a glance. We recognize, however, that this may have led to some difficulty in comparing specific related groups due to spatial constraints. We sincerely apologize for any confusion and are grateful for the reviewer's suggestion of selectively highlighting intriguing groups to emphasize commonalities more clearly. We believe this will markedly enhance the interpretability of our findings. Accordingly, we will revise the figure as Figure R1 in the PDF attached to the global response.\n\n- Generated Imaginary Tasks: We acknowledge that basis subtasks may not always align with human intuition, and it may be challenging to assert that a specific dimension corresponds to a particular action. The essence is to construct augmented imaginary tasks with various compositions, each exhibiting distinct behaviors. For example, one task might ignore an object, while another may involve picking it up. To better interpret these varied behaviors, please refer to our anonymous web pages (linked in Line 305 of the main manuscript) for the videos of generated virtual tasks that had been uploaded before the time of the manuscript submission.\n\n### **W5.** Minor Issues\n\n- Dropout: We will provide further clarification that dropout is applied to the embeddings of states and actions, following the same method as LDM [2].\n\n- Figures 2 and 3: To clarify, Figure 2 pertains to the training tasks, while Figure 3 is designed to illustrate the generation of the imaginary reward, denoted as $\\tilde{r}$ given a fixed imaginary composition $\\tilde{y}$.\n\n- Use of \"zero-shot\": We recognize that our use of the term \"zero-shot\" to describe the absence of gradient updates may have led to confusion. To ensure clarity, we will revise this description in our manuscript.\n\n### **Q2.** Overfitting to High Reward States\n\nThe decoder is optimized over the distribution of states and actions that the policy visits. As training progresses, the policy tends to exploit states and actions associated with higher rewards, causing the decoder to overfit these states and actions. This phenomenon aligns with the observation in Figure 4 of the LDM paper [2], illustrated with a Gridworld example. This is precisely where dropout comes into play. It ensures that the decoder doesn't solely rely on the states and actions but also takes the context into account when predicting the reward.\n\n### **Q3.** Encoding Sequential Subtasks\n\nIn our current model, the sequence or order of subtasks is not explicitly encoded. We hypothesize that this information is implicitly captured within the continuous contexts of the model. Enhancing the model to deliberately include the sequential order of subtasks is an intriguing direction for future research and could provide additional insights.\n\n### **Q4.** Limitations of LDM for Parametric Variants\n\nLDM produces virtual tasks through interpolations of contexts from a single Gaussian VAE. By referring to Figure 9 of the LDM paper [2], it can be inferred that such contexts capture the parametric details of the tasks (e.g., coordinates of goals). In contrast, our method conducts virtual training in the subtask space, allowing for the generation of unseen compositional tasks."}, {"review_id": "CU4kxYgeU1", "reviewer": "Reviewer_4fbE", "summary": "The 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.", "questions": "1. 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", "limitations": "The authors addressed the limitations of the paper", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for recognizing our contribution and posing many insightful questions. We are confident that by addressing these points, our manuscript will offer a more comprehensive and clearer understanding.\n\n### **W1.** and **Q1.** \"Non-parametric\" Problem Formulation\n\nWe apologize for the confusion, we will revise it as our global response **G1.**\n\n### **W3.** Simple MuJoCo Tasks\n\nWe evaluated SDVT-LW on 3 MuJoCo tasks defined in Table 1 of the LDM paper[2]. We report the result in Figure R2 attached to the global response. Ours performs comparable to baselines although it's not specifically designed nor tuned for simple parametric variations. We used the same hyperparameters as LDM, except for components unique to ours, such as the GMVAE: $K=4$ for Ant-direction, and $K=2$ for Ant-goal and Half-cheetah-velocity. We referred to the results of the baselines reported in Figure 7 of the LDM paper. \n\n### **W4.** Capacity Matching\n\nWe agree that we should match the capacity component-wise for a fair comparison. We have made adjustments to VariBAD and LDM accordingly: added an MLP with hidden sizes [1600, 256] after $h\\_t$ in the encoder (post-GRU output), added a final layer of size [128] to the policy, and enlarged the decoder's hidden dimensions from [64,64,32] to [128,256,160]. The matched VariBAD components have slightly more parameters than those of SDVT-LW. Table R1 in the attached PDF indicates that while the matched baselines improve, they still lag behind ours.\n\n### **Q2.** Occupancy Regularization\n\nWe address the occupancy regularization in our global response **G2.**. In addition to what is discussed there, it's important to note that the occupancy regularization is not solely intended to influence the virtual training. Rather, it compresses the effective dimension of subtasks to be generalized for the test. We run the same ablation for SD with $\\alpha\\_o=0.0, 5.0, 10.0$ that scored test success rates of 22.8\\%, 28.4\\%, and 11.2\\%, respectively on ML-10. The default SD with $\\alpha\\_o=1.0$ scored 30.8\\%.\n\n### **Q3.** Multivariate Gaussian\n\nWe apologize for any confusion from the lack of detailed explanations. To clarify, both VariBAD and LDM employ a single multivariate Gaussian encoder to capture contexts. Our approach formulates the encoder as a mixture of such multivariate Gaussian encoders. It's worth noting that there are multiple ways to configure a Gaussian mixture. Our implementation differs from those marginalizing the categorical variable $y\\_{t}$, but is conceptually similar to the M2 model presented in [4]. As the reviewer rightly pointed out, the input to the mean and variance MLPs is the concatenation of $h\\_t$ and $y\\_t$. For a more comprehensive understanding of our formulation, we direct the reviewer to the GMVAE repository we linked in Appendix C.1. (Line 451). We intend to incorporate these clarifications into the final manuscript to enhance comprehension.\n\n### **Q4.** Dropout\n\nYes, we'll state that the dropout structure is exactly the same as LDM. \n\n### **Q5.** Less Rollouts\n\nWe use $n\\_{\\text{roll}}=10$, following the original setup of the Meta-World benchmark [1]. Notably, this benchmark utilizes dense rewards, so the context converges rather quickly, even within the first rollout. This rapid convergence was illustrated in Figure 2(b) and our supplementary page (Line 305). As a result, the performance across rollouts does not exhibit substantial improvement. To provide further insight, we report the success rates for smaller $n\\_{\\text{roll}}$ in Table R2 of the attached PDF.\n\n### **Q6.** Sampling Imaginary Composition\n\nIn the case of LDM, a weight vector $(\\alpha^{(1)},\\ldots, \\alpha^{(n)})$ is sampled from a Dirichlet($\\mathbf{1}$) with all-one parameters and is used to compute the weighted sum of parametric contexts from $n$ independent workers from training tasks, given by $\\tilde{z}\\_{t}=\\sum\\_{i=1}^{n} \\alpha^{(i)} z^{(i)}\\_{t}$. Whereas we sample the composition, $\\tilde{y} \\sim$ Dirichlet($\\bar{y}$) with the mean $\\bar{y}$ being the running mean of training task compositions.\n\nThe rationale behind our choice is that $\\tilde{z}\\_{t}$ can be any continuous value, whereas $\\tilde{y}$ must be a probability vector. Therefore, the Dirichlet distribution seemed appropriate to generate unseen but similar compositions. Also using $\\bar{y}$ instead of $\\mathbf{1}$ as the distribution's parameter performed more stably. However, we acknowledge that this choice does not have a robust theoretical standing, and we are open to exploring alternatives. The reviewer's suggestion to consider using other theoretically solid approaches, such as KDE, for sampling virtual contexts could be an exciting avenue for future work.\n\n### **Q7.** Virtual Training on Imaginary Contexts\n\nOur approach is similar to LDM which also uses a time-varying mixture context, $\\hat{m}\\_t$ (eq.(1) of the LDM paper), corresponding to $\\tilde{z}\\_t$ in our case. To replicate the test, it seems intuitive to update $\\tilde{z}\\_t$ over time since it's a function of $h\\_t$, which evolves in response to the previously generated virtual reward, $\\tilde{r}\\_t$. In simpler terms, maintaining $\\tilde{y}$ constant is akin to preserving the $\\alpha^{(i)}$s in LDM during the virtual meta-episode, yet it's important to note that the final contexts are not kept static.\n\n### **Q8.** Finite Task Space\n\nUnlike some prior works on MuJoCo, which consider infinite parametric variations, the Meta-World benchmark operates under a finite task space. This constraint adds to the challenge of generalization but also contributes to a more stable evaluation. In the case of the ML-10 example, upon initializing the benchmark, all 10 training and 5 test tasks sample 50 parametric variants each, based on the random seed. This pool of 15$\\times$50 tasks remains fixed throughout the entire training process.\n\n### **Q9.** Related Work\n\nWe thank the reviewer for additional related works that deserve mention in our paper."}, {"review_id": "CQcqWxnn62", "reviewer": "Reviewer_PDTa", "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.", "questions": "1. Why SDVT does not use imagined dynamics for training? Not using imagined dynamics may decease SDVT's generalization power.\n2. 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]?\n\n\n\n[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?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Dealing with multi-modal and diverse task distributions is an important problem in Meta-RL.\n2. Presentation is clear, and the paper is easy to follow.\n3. The methodology is sensible.\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for acknowledging the importance of the problem and the methodology. We address the questions as follows.\n\n### **Q1.** Why not Imagine State Transition Dynamics?\n\nOur current approach focuses on tasks with varying reward dynamics and doesn't extend to imagining unseen state-transition dynamics. When attempting to use an MLP state-transition decoder to create imaginary tasks, we encountered significant errors that accumulated over longer timesteps, such as in the Meta-World benchmark (5000 steps per meta-episode). Employing an advanced generative model, such as the Transformer mentioned in **Q2.**, could potentially allow for more accurate generation of long state transitions. This remains a promising avenue for future research within our RL community.\n\n### **Q2.** Comparison to a Related Work\n\nThank you for introducing the interesting related work, MILLION [5] that merits a reference in our paper. Our method differs significantly from MILLION in its focus and approach. While MILLION emphasizes language instructions and utilizes a Transformer encoder, our method employs a GRU encoder and specifically leverages a Gaussian mixture VAE to model the latent space of compositional tasks.\n\nThe advantage of our method lies in the explicit use of mixture components in latent features, while its limitation stems from the restricted expressive power of the GRU compared to the Transformer. We chose the GRU over the Transformer because of computational considerations; the computation required for the Transformer grows quadratically with input size, while the GRU's computation grows linearly. In a benchmark like Meta-World, with inferences over very long horizons (up to 5000 steps), the GRU was more suitable.\n\nNevertheless, we recognize that the use of a Transformer might eventually be necessary for both effective inference and the generation of virtual dynamics (as discussed in **Q1.**). Exploring ways to compress long-horizon meta-RL tasks using a Transformer could be a promising direction for future research.\n\n### **Q3.** Computation\n\nWe direct the reviewer to Appendix C.2, where the relative computational complexity is detailed. The increase in computational burden in our method is from incorporating the GMVAE over a simple VAE, as well as the generation of virtual tasks. While these additional complexities do raise the computational requirements, the demand is still substantially within the same range as those of VariBAD and LDM.\n"}]

Author rebuttal:
We appreciate the reviewers for dedicating their time and effort to recognize the strengths of our paper, and for offering constructive feedback that will undoubtedly enhance the presentation of our work. 

Below are global responses that address concerns raised by multiple reviewers. Additionally, we have provided separate, detailed responses tailored to the comments and suggestions of each individual reviewer. We report additional results in 2 figures and 2 tables in the attached PDF file.

### **G1.** Clarification of the Term "Non-Parametric" Task Generalization

The terminology "non-parametric" in relation to task variability was first coined in the Meta-World paper [1]. Its purpose was to differentiate the task variability in Meta-World manipulation tasks from that in the standard MuJoCo tasks, which primarily involve simple parametric variations such as target goal positions, directions, and velocities.

We acknowledge, as underscored by Reviewers juhC and 4fbE, that those "non-parametric" tasks can indeed be parametrized, in terms of the subtask compositions (or task index) and the traditional parametrization of each subtask. A critical insight of our work is the recognition that such specific parametrization was challenging to achieve with prior methods that utilized a single Gaussian VAE. Our primary innovation lies in facilitating this type of parametrization through the application of a Gaussian mixture VAE, which allows for a more appropriate understanding and handling of such tasks.

In conclusion, we acknowledge that the term "non-parametric" could lead to misunderstandings, as it may imply an absence of parameters to parametrize the tasks, whereas our approach successfully parametrizes them. Therefore, in order to mitigate any confusion, we will clarify the scope of our work by referring to it as modularity/composition-based generalization, in line with the suggestions made by the reviewers.

### **G2.** Occupancy Regularization (SDVT-LW vs SDVT)

We emphasize the critical role of the number of subtasks (Line 160), denoted by $K$, as illustrated in Table 10 (Appendix E.2). Since the optimal value for $K$ may not always be accessible, we propose using occupancy regularization as a remedy in such cases. 

As the reviewers suggested, we can initiate $K$ with an adequately large value such as $N\_{\textrm{train}}$. However, as outlined in the paragraph beginning from Line 160, this approach has a susceptibility: each task may simply be classified into one-hot subtask composition, rather than a mixture of subtasks. While this might enhance training performance, it becomes problematic during testing with unseen compositions of subtasks. In these circumstances, the parametrization struggles to generalize, as evidenced by the low test success rate for $K=10$ in Table 10 (Appendix E.2). 

To mitigate this, we employ occupancy regularization, intentionally suppressing higher dimensions to prevent the subtask representation from being reduced to one-hot labels. Our response to **W1.** of Reviewer juhC further elaborates on this approach and may help alleviate this concern. The empirical gain of the occupancy regularization is reported in Appendix E.1. In essence, while SDVT-LW is built on the presumption of knowing the optimal $K$, SDVT extends SDVT-LW by integrating occupancy regularization, enabling it to navigate broader scenarios without specific knowledge of the ideal $K$.
  
**References.** The list of related works referred to in our responses is as follows

[1] Yu et al. "Meta-World: A benchmark and evaluation for multi-task and meta reinforcement learning." arXiv preprint 2021 (an updated version of CoRL 2019).

[2] Lee and Chung. "Improving generalization in meta-RL with imaginary tasks from latent dynamics mixture." NeurIPS 2021.

[3] Ajay et al. "Distributionally adaptive meta reinforcement learning." NeurIPS 2022.

[4] Kingma et al. "Semi-supervised learning with deep generative models." NeurIPS 2014.

[5] Bing et al. "Meta-reinforcement learning via language instructions." arXiv preprint 2022.

Meta-review: This paper proposes a new architecture for generalization to RL tasks with non-parametric variation, where the main idea is train a Gaussian mixture VAE to meta-learn the subtask decomposition process. All the reviewers agreed that the proposed method is well-motivated and sufficiently novel for the given problem, and the results on Meta-World are strong and convincing. The authors clarified most of the reviewers' questions and added more results, which addressed some of the reviewer's concerns about unfair comparisons and weak analysis. Thus, I recommend accepting this paper.
Decision: Accept (poster)