SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The overall objective in Eq. (6) is optimized over the trajectories of training tasks. To enhance generalization to test tasks with unseen subtask compositions, we generate virtual training tasks using the imaginary dynamics produced by the GMVAE decoder. This process resembles those in [35, 1], but their generated tasks are limited to parametric variations, e.g., generating tasks with unseen goal positions. Such methods can prepare for test tasks with unseen parametric changes but struggle to prepare for qualitatively new tasks with unseen compositions of subtasks. Our GMVAE model enables us to extend the process to the non-parametric setup by conditioning the decoder on an imaginary subtask composition  $\tilde{y}$ . At the beginning of each meta-episode, we randomly determine with probability  $p_v$  whether to convert it into a virtual meta-episode. By training the policy on imaginary tasks, it can better prepare for test tasks with unseen subtask compositions in advance. We use the tilde accent to denote imaginary components, and the hat accent to denote estimates.
Latent context dispersion Inspired by the work of Ajay et al. [1], we utilize the dispersion structure for our GMVAE to support extrapolated generalization of the latent context  $z_{t}$ . Instead of directly using  $z_{t}$  as the decoder input, we insert an additional MLP  $p_{\theta_{\hat{h}}}$  before the decoder to expand the dimension of the context (the dotted box in Figure 2a). The MLP output  $\hat{h}_t = p_{\theta_{\hat{h}}}(z_t)$  is trained to reconstruct the embedding  $h_t$  by appending  $\alpha_d\mathcal{I}_{\mathrm{dis}} = -\alpha_d\| h_t - \hat{h}_t\| _2^2$  to the total GMVAE objective in Eq. (3). We then use the dispersed context  $\hat{h}_t$  instead of  $z_{t}$  to reconstruct the reward and transition. This trick is effective in generating imaginary tasks featuring extrapolated dynamics, albeit at the cost of increased training complexity.
Imaginary reward generation Figure 3a presents the imaginary reward generation process. The goal is to create imaginary tasks based on the distribution of training subtasks but with unseen compositions of subtasks. At the beginning of a virtual meta-episode, we randomly sample an imaginary subtask composition  $\tilde{y} \sim \mathrm{Dirichlet}(\bar{y})$  fixed for that virtual meta-episode, where the concentration parameter  $\bar{y} \in \Delta^K$  is the empirical running mean of all  $y_t$  over training. By replacing the real  $y_t$  with the imaginary subtask composition  $\tilde{y}$ , we sample an imaginary context  $\tilde{z}_t \sim \mathcal{N}(\mu_{\phi_z}(h_t, \tilde{y}), \sigma_{\phi_z}^2(h_t, \tilde{y}))$ , imaginary dispersed context  $\hat{\tilde{h}}_t = p_{\theta_{\hat{h}}}(\tilde{z}_t)$ , and finally the imagi
nary reward  $\tilde{r}_{t + 1}\sim p_{\theta_R}(r_{t + 1}|\mathrm{DO}[s_t,a_t,s_{t + 1}],\tilde{\hat{h}}_t)$  accordingly using our GMVAE. We replace the reward for the next timestep,  $r_{t + 1}$ , with  $\tilde{r}_{t + 1}$ , while the states remain to be from the real training task. The imaginary reward is used for the encoder input at the next time step and for the policy, where the policy is trained to maximize the sum of generated rewards,  $\sum_{t = 0}^{H^{-1}}\gamma^t\tilde{r}_{t + 1}$ . However, the GMVAE is not trained to reconstruct the imaginary dynamics.
3.3 Summary of the Combined Methods: SDVT-LW and SDVT
By combining the subtask decomposition (SD) and virtual training (VT) processes, we propose two methods: SDVT-LW and SDVT. Foremost, our primary contribution is the proposal of SDVT-LW, which is the lightweight (-LW) version of our method that assumes the prior knowledge of the optimal number of subtasks  $K$ , therefore not employing the occupancy regularization.
Furthermore, we propose SDVT with the occupancy regularization strategy. This generalizes 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 [1]. Their use in our work remains unchanged from their original applications. The efficacy of these components, within the context of virtual training, is demonstrated via ablations in Appendix E. We summarize the entire meta-training process with a pseudocode in Appendix A.
Classical meta-RL Classical meta-RL methods assume that both training and test tasks are sampled from the same distribution. These methods are divided into two main categories: gradient-based methods and context-based methods. Gradient-based methods [13, 54, 48, 73, 53] learn a common initialization of a lightweight model for all tasks, allowing the agent to achieve high performance on unseen target tasks with a few steps of gradient updates. However, these methods lack online adaptation capability within a meta-episode because they require many pre-update rollouts before adaptation. Context-based methods [9, 18, 46, 32, 34, 43, 74, 40] use a recurrent or memory-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 out-of-distribution tasks.
Out-of-distribution meta-RL A group of recent studies focuses on training a generalizable agent that is robust to test-time distribution shifts. A group of works generates imaginary observations using image augmentation techniques [18, 21, 31, 44, 33, 61, 66]. Most of these methods depend on predefined heuristic augmentations, without utilizing the training task dynamics. On the other hand, some works explicitly address varying environment dynamics. For example, MIER [7] reuses the trajectories collected during training by relabeling according to the test dynamics. Our work is related to AdMRL [36], LDM [35], and DiAMetR [1], which generate imaginary tasks with unseen dynamics using learned models. However, these works focus on generating parametric variants of training tasks, while we focus on generalizing across non-parametric task variants.
Skill-based meta-RL Our task parameterization based on the subtask decomposition is related to the recently spotlighted skill-based meta-RL methods [11, 51, 50, 55], which aim to achieve fast generalization on unseen tasks by decomposing and extracting reusable skills from training tasks' trajectories. These works often require refined offline demonstrations to learn skills using behavioral cloning objectives, where the skills distinguish a sequence of actions given a sequence of states. For example, SimPL [42] extracts skills from offline demonstrations before meta-training, and during a meta-test, it only adapts the high-level skill policy, with the low-level policy frozen. HeLMS [47] learns a 3-level skill hierarchy by decomposing offline demonstrations of a single task. Online learning of the skills is often unstable because the set of skills should develop along with the online improvement of the policy. The subtask decomposition of our method is conceptually different from
the skill decomposition [68]. Even when the policy is not stationary during online updates, the underlying reward and transition dynamics, which our model has to estimate, do not change.
Meta-World benchmark The Meta-World V2 benchmark [71] stands as the most prominent, if not the only, established benchmark for assessing meta-RL algorithms featuring non-parametric task variability. This benchmark comprises 50 qualitatively distinct robotic manipulation tasks, with each task containing 50 parametric variants that incorporate randomized goals and initial object positions. Specifically, the Meta-World Meta-Learning 10 (ML-10) benchmark, consists of  $N_{\mathrm{train}} = 10$  training tasks and  $N_{\mathrm{test}} = 5$  held-out test tasks. We denote each task by an index, where the training tasks are numbered from 1 to 10 and the test tasks from 11 to 15. Likewise, the ML-45 benchmark consists of  $N_{\mathrm{train}} = 45$  training tasks and  $N_{\mathrm{test}} = 5$  test tasks. Refer to the tables in Appendix D.1 for the set and indices of tasks. The agent must maximize its return from experience while exploring to identify the initially unknown task dynamics within a meta-episode of  $H^{+} = 5000$  steps that consists of  $n_{\mathrm{roll}} = 10$  rollout episodes of horizon  $H = 500$  steps each.
SDVT variants and baselines setup We evaluate our methods SDVT and SD (only subtask decomposition without virtual training and dispersion). Without the prior knowledge of  $K^*$ , we set  $K = N_{\mathrm{train}}$  and apply the occupancy regularization  $(\alpha_o = 1.0)$  with  $\alpha_c = 1.0$ . We also evaluate a lightweight (-LW) version of ours with smaller  $K = 5$  with  $\alpha_c = 0.5$  and without the occupancy regularization  $(\alpha_o = 0.0)$ . To ensure a fair comparison and to exclude the gains from orthogonal contributions, we compare SDVT with state-of-the-art meta-RL methods that do not require any refined offline demonstrations, ensembles, or extensive test-time training:  $\mathrm{RL}^2$  [9], MAML [13], PEARL [46], and VariBAD [74]. We also compare with a parametric OOD meta-RL method, LDM [35], to evaluate the efficacy of our virtual training over subtasks. All methods do not perform gradient updates during the test except for MAML. Appendix C presents more implementation details. Briefly, SDVT without a Gaussian mixture reduces to LDM, LDM without virtual training reduces to VariBAD, and VariBAD without a VAE decoder reduces to  $\mathrm{RL}^2$ . Our implementation is available at https://github.com/suyoung-lee/SDVT.
Evaluation metric We follow the standard success criterion of Meta-World as follows. A timestep is considered successful if the distance between the task-relevant object and the goal is less than a predefined threshold. A rollout episode is considered successful if the agent ever succeeds at any timestep during the rollout episode. The success of a meta-episode is defined as the success of the last (10th) rollout episode. In Table 1, we report the success rates of 750 (15 tasks  $\times$  50 parametric variants) meta-episodes at 250M steps for ML-10 and 2500 (50 tasks  $\times$  50 parametric variants) meta-episodes at 400M steps for ML-45.2 Likewise, we report the returns of the last rollout episodes.
Table 1 illustrates that no method attains a  $100\%$  success rate even on training tasks, emphasizing the challenge posed by non-parametric task variability. The training success rates on ML-45 are consistently lower than those on ML-10, reflecting the inherent difficulty in adapting to a broader range of tasks, such as conflicting gradients [70]. Notably, SD surpasses all other baselines in training success rate. In particular, outperforming VariBAD underscores the limitations of employing a single Gaussian distribution to model the latent task space in cases of non-parametric variability.
On the test tasks, SDVT and SDVT-LW substantially outperform all baselines, even outperforming LDM, which surpasses VariBAD with parametric virtual training. Our gain is attributed to our virtual training process, which is specifically designed for test tasks involving non-parametric variability. Notably, SDVT and SD outperform their LW counterparts in training success rates, primarily due to its fine-grained subtask decomposition which provides a more precise representation of each task. For example, a "grip object" subtask may be split and represented as a combination of two distinct subtasks "move gripper" and "tighten gripper" with a larger  $K$ . In contrast, SDVT-LW scores higher
2Our aggregation method diverges from the atypical method employed in the Meta-World paper [71], which presents the main results as the average of the maximum success rate of each task (details in Appendix C.5)
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 and ML-45 benchmarks. All results are reported as the mean success rate  $\pm 95\%$  confidence interval of 8 seeds. Individual scores of all tasks are reported in Appendix D.1.
Methods	Success Rate	Return
ML-10	ML-45	ML-10	ML-45
Train	Test	Train	Test	Train	Test	Train	Test
SDVT	77.2±3.0	32.8±3.9	55.6±4.2	28.1±3.2	3656±62	1225±160	2379±214	839±74
SDVT-LW	62.1±4.1	33.4±5.0	50.4±4.1	31.2±1.2	3454±137	1527±214	2294±202	894±27
SD	77.0±5.9	30.8±7.7	61.0±1.7	23.0±5.1	3630±241	1112±190	2672±79	786±69
SD-LW	75.5±5.5	26.2±8.7	56.7±1.5	25.4±2.9	3525±297	1043±234	2578±64	793±49
RL²	67.4±4.4	15.1±2.7	58.0±0.4	11.8±3.2	1159±83	715±33	1411±22	663±100
MAML	42.2±4.5	3.9±3.7	32.0±1.4	19.8±6.3	1822±136	439±78	1388±104	658±96
PEARL	23.2±1.9	0.8±0.5	10.3±2.4	6.7±3.3	1081±77	340±54	597±121	506±122
VariBAD	58.2±8.9	14.1±6.1	57.0±1.2	22.1±3.5	3055±466	919±143	2492±47	762±40
LDM	56.7±12.3	19.8±6.0	54.1±0.9	24.8±2.9	2963±626	1166±264	2515±67	768±63
test success rates than SDVT, presumably due to its coarser decomposition that allows imaginary compositions to encompass a broader range of unseen test tasks.
Please refer to the tables in Appendix D.1 for individual task results. Our methods achieve the highest success rates across all test tasks on the ML-10 benchmark. However, all methods encounter challenges in solving "Shelf-place," which includes an unseen shelf not present in the observation. As such, this task cannot be decomposed into previously seen subtasks but rather into unseen ones, making it difficult to prepare through virtual training. These tasks, composed of unseen subtasks, pose a considerable challenge for zero-shot adaptation as with SDVT. Addressing these tasks may require a substantial increase in rollouts and updates during tests, an area of potential future work. Detailed ablation results are reported in Appendix E. For additional results such as rendered videos, refer to our webpage https://sites.google.com/view/sdvt-neurips.
In this subsection, we explore whether the performance enhancements attributed to our methods are indeed the result of effectively decomposed subtasks and a diverse range of imaginary tasks.
Learned Subtask Compositions From Figure 2b, we observe a rapid context convergence within the first rollout episode for a sufficiently meta-learned task. To validate our motivation, we visualize the converged subtask compositions in Figure 4. We find that such converged subtask compositions are shared by qualitatively similar tasks. For example, in Figure 4a, tasks (3) "Pick-place," (7) "Peg-insert-side," and (10) "Basketball," which require placing an object at a goal position, share subtask indices 1 and 3. Additionally, simple tasks that require moving the gripper in a straight line: (1) "Reach," (5) "Drawer-close," and (8) "Window-open" are classified into the same subtask 2. These observations suggest that our model can effectively meta-learn the decomposition process of tasks into shareable subtasks. Furthermore, the test tasks may not necessarily have the same subtask compositions as the training tasks. For instance, (3) "Pick-place" and (14) "Sweep-into" share subtask 3, but not subtasks 1 or 5, revealing the potential for virtual training to be effective.
Occupancy Regularization Figure 4 indicates that it is important to set the subtask class dimension  $K$  and the categorical coefficient  $\alpha_{c}$  appropriately according to the number of training tasks and the correlations among them. Otherwise, the decomposition may suffer from a collapse or degeneration. With the occupancy regularization, we can avoid the extensive search for the optimal subtask dimension  $K^{*}$  and corresponding  $\alpha_{c}$ . In Figure 4b, we find that our occupancy regularization successfully limits the use of unnecessary subtasks with higher indices as intended.
Generated imaginary tasks To demonstrate the dynamics of the imaginary tasks, we show the trajectory of the gripper and object over a meta-episode (5000 steps) on generated imaginary tasks for different imaginary subtask compositions  $\tilde{y}$  in Figure 3. When we set  $\tilde{y}$  as a one-hot vector, the policy
(a) SDVT-LW  $(K = 5, \alpha_{c} = 0.5, \alpha_{o} = 0.0)$ .
(b) SDVT  $(K = 10, \alpha_{c} = 1.0, \alpha_{o} = 1.0)$ .
(c) SDVT-LW  $(K = 5,\alpha_{c} = 0.1,\alpha_{o} = 0.0)$
(d) SDVT-LW  $(K = 5,\alpha_{c} = 2.0,\alpha_{o} = 0.0)$
Figure 4: Learned subtask compositions on ML-10. (a) Default SDVT-LW (b) Default SDVT. (c) and (d) SDVT-LW with varying  $\alpha_{c}$ . Each column displays the terminal subtask composition ( $y_{H^{+}}$ ) of each task learned after 250M training steps of training, averaged across 50 parametric variants. The decompositions at different timesteps can be found in 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.
returns a trajectory that solves an elementary subtask, such as placing and reaching up. We observe that the trajectories vary across different  $\tilde{y}$ , and similar compositions result in similar trajectories.
In conclusion, our proposed method has demonstrated a considerable enhancement in meta-RL with non-parametric task variability. This improvement is achieved by meta-learning shareable subtask decompositions and executing virtual training on the imaginary subtask compositions. However, it's essential to acknowledge certain potential limitations beyond the scope of this study, particularly when addressing test tasks involving entirely novel subtasks and in broader setups where the action and state spaces may also vary. Despite these limitations, we believe that expanding the realm of meta-RL to accommodate a wider range of task variability is a critical research topic. Incorporating orthogonal approaches such as using offline demonstrations or test time training techniques into our method could lead to interesting future work addressing the limitations. We are optimistic that our study lays a robust foundation for future research in this field.
This work was supported in part by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2022-0-00469, Development of Core Technologies for Task-oriented Reinforcement Learning for Commercialization of Autonomous Drones,  $50\%$  ) and in part by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2022-0-00124, Development of Artificial Intelligence Technology for Self-Improving Competency-Aware Learning Capabilities,  $50\%$  )
Algorithm 1 Subtask Decomposition and Virtual Training (SDVT)
Initialize Encoder  $q_{\phi}$  decoder  $p_{\theta}$  policy  $\pi_{\psi}$  , set of training tasks  $\mathcal{M}_{\mathrm{train}}$  , virtual ratio  $p_v$  , GMVAE buffer  $\mathcal{B}_{\mathrm{VAE}}$  , policy buffer  $\mathcal{B}_{\mathrm{pol}}$  , total number of meta-episodes to train on  $n_{\mathrm{meta}}$  , number of rollout episodes per meta-episode  $n_{\mathrm{roll}}$  , running mean of subtask composition  $\bar{y}$    
for meta-episode  $k = 0,\dots ,n_{\mathrm{meta}} - 1$  do Sample a training task  $M_{k}\sim p(\mathcal{M}_{\mathrm{train}})$  Sample a virtual meta-episode flag  $V\sim$  Bernoulli  $(p_v)$  if  $V = 1$  then Sample an imaginary subtask composition  $\tilde{y}\sim$  Dirichlet(y)   
end if Reset  $h_0,y_0,\mathcal{B}_{\mathrm{pol}}$    
for timestep  $t = 0,\ldots ,n_{\mathrm{roll}}\times H - 1$  do if  $t$  mod  $H = 0$  then Reset rollout episode  $s_t\sim T_{0,k}(\cdot)$    
end if Sample an action  $a_{t}\sim \pi_{\psi}(\cdot |s_{t},\mu_{\phi z}(h_{t},y_{t}),\sigma_{\phi z}(h_{t},y_{t}),\omega_{\phi y}(h_{t}))$  Take an environment step  $s_{t + 1}\sim T_k(\cdot |s_t,a_t)$ $r_{t + 1}\gets R_k(s_t,a_t,s_{t + 1})$    
if  $V = 1$  then  $\tilde{z}_t\sim \mathcal{N}(\mu_{\phi_z}(h_t,\tilde{y}),\sigma_{\phi_z}^2 (h_t,\tilde{y}))$ $\tilde{\hat{h}}_t = p_{\theta_{\hat{h}}}(\tilde{z}_t)$ $\tilde{r}_{t + 1}\sim p_{\theta_R}(\cdot |\mathrm{DO}[s_t,a_t,s_{t + 1}],\tilde{\hat{h}}_t)$  Replace  $r_{t + 1}\leftarrow \tilde{r}_{t + 1}$    
else Add the transition  $(s_t,a_t,s_{t + 1},r_{t + 1})$  to  $\mathcal{B}_{\mathrm{VAE}}\dots$  VAE NOT TRAINED WITH VIRTUAL DYNAMICS   
end if Add the transition  $(s_t,a_t,s_{t + 1},r_{t + 1})$  to  $\mathcal{B}_{\mathrm{pol}}$  Update hidden embedding  $h_{t + 1} = q_{\phi_h}(\tau ;t + 1)$  Update subtask composition  $y_{t + 1}\sim$  Cat  $(\omega_{\phi_y}(h_{t + 1}))$  Update the running mean of subtask composition  $\bar{y}\gets$  RunningMeanUpdate(  $\bar{y},y_{t + 1}$    
end for Update GMVAE  $\{\phi ,\theta \} \leftarrow \{\phi ,\theta \} +\nabla_{\{\phi ,\theta \}}\sum_{t = 0}^{H^+}\mathrm{ELBO}_t$  with samples from  $\mathcal{B}_{\mathrm{VAE}}$  Update policy  $\psi \leftarrow \psi +\nabla_{\psi}\mathcal{I}_{\mathrm{pol}}^{+}$  with samples from  $\mathcal{B}_{\mathrm{pol}}$  Anneal virtual ratio  $p_v\gets p_v + \Delta p_v$    
end for
B ELBO Derivation
The GMVAE's ELBO objective is derived as follows.
$$
\begin{array}{l} \mathbb {E} _ {d (M _ {k}, \tau_ {: H ^ {+}})} \left[ \log p _ {\theta} (\tau_ {: H ^ {+}}) \right] = \mathbb {E} _ {d (M _ {k}, \tau_ {: H ^ {+}})} \left[ \log \mathbb {E} _ {q _ {\phi} (y _ {t}, z _ {t} | h _ {t})} \left[ \frac {p _ {\theta} (\tau_ {: H ^ {+}} , y _ {t} , z _ {t})}{q _ {\phi} (y _ {t} , z _ {t} | h _ {t})} \right] \right] \\ \geq \mathbb {E} _ {d (M _ {k}, \tau_ {: H ^ {+}})} \left[ \mathbb {E} _ {q _ {\phi} (y _ {t}, z _ {t} | h _ {t})} \left[ \log \frac {p _ {\theta} (\tau_ {: H ^ {+}} , y _ {t} , z _ {t})}{q _ {\phi} (y _ {t} , z _ {t} | h _ {t})} \right] \right] \\ = \mathbb {E} _ {d \left(M _ {k}, \tau_ {: H ^ {+}}\right)} \left[ \mathbb {E} _ {q _ {\phi} \left(y _ {t}, z _ {t} \mid h _ {t}\right)} \left[ \log \frac {p _ {\theta} \left(\tau_ {: H ^ {+}} \mid y _ {t} , z _ {t}\right) p _ {\theta} \left(z _ {t} \mid y _ {t}\right) p \left(y _ {t}\right)}{q _ {\phi} \left(y _ {t} \mid h _ {t}\right) q _ {\phi} \left(z _ {t} \mid h _ {t} , y _ {t}\right)} \right] \right] \\ = \mathbb {E} _ {d \left(M _ {k}, \tau_ {: H ^ {+}}\right)} \left[ \mathbb {E} _ {q _ {\phi} \left(y _ {t}, z _ {t} \mid h _ {t}\right)} \left[ \log p _ {\theta} \left(\tau_ {: H ^ {+}} \mid z _ {t}\right) + \log \frac {p _ {\theta} \left(z _ {t} \mid y _ {t}\right)}{q _ {\phi} \left(z _ {t} \mid h _ {t} , y _ {t}\right)} + \log \frac {p \left(y _ {t}\right)}{q _ {\phi} \left(y _ {t} \mid h _ {t}\right)} \right] \right]. \tag {8} \\ \end{array}
$$
Eq. (8) is equivalent to the ELBO objective in Eq. (3) without weighting coefficients. We assume that the reconstruction  $\tau_{:H^{+}}$  is conditionally independent of the subtask composition  $y_{t}$  given  $z_{t}$ .
C.1 Reference Implementations
MAML, RL², and PEARL To replicate the results of MAML [13], RL² [9], and PEARL [46] as reported in the Meta-World paper [71], 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.