SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Suyoung Lee, Myungsik Cho, Youngchul Sung*
Daejeon 34141, Republic of Korea
{suyoung.l, ms.cho, ycsung}@kaist.ac.kr
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.
Meta-reinforcement learning (meta-RL) constitutes a dynamic field within deep reinforcement learning, focusing on training agents to quickly adapt to novel tasks by learning from a variety of training tasks [2]. By interacting with these tasks, meta-RL creates an inductive bias regarding the task dynamics and subsequently develops a policy based on this knowledge. Despite its significant contribution to the generalization capability of traditional deep RL, meta-RL is susceptible to test-time distribution shifts, which restricts its applicability to familiar in-distribution test tasks [17, 36, 39, 41].
To tackle this limitation, recent out-of-distribution (OOD) meta-RL approaches have emphasized distinct training and test task distributions, thereby achieving enhanced performance on unseen OOD test tasks with interpolated or slightly extrapolated training dynamics [12, 36, 41, 35, 1]. Although the parameters of training and test tasks are drawn from disjoint distributions, these tasks remain qualitatively similar, as they can be expressed in a shared parametric form representing the task dynamics (e.g., the same "Pick-place" task with OOD goal positions in Figure 1b).
In this study, we explore a more general meta-RL framework that addresses non-parametric task variability [69, 71], a topic that has received limited attention in prior research. In this context as of Figure 1c, variations among tasks cannot be expressed through simple parametric variations, such as the parameterization of a goal position. Generalization is particularly challenging in this setting, as conventional meta-RL methods often model the inductive bias as a parametric embedding applicable to various tasks [9, 46, 74]. Within a non-parametric framework, it may not be feasible to employ a unified and generalizable parameterization of training tasks using standard meta-RL techniques.
(a) "Pick-place" in standard indistribution meta-RL setup.
(b) "Pick-place" in out-of-distribution meta-RL setup.
(c) Non-parametric task variation between "Pick-place" and "Sweep-into."
Figure 1: Problem Setup. Visualizing different meta-RL scenarios with Meta-World tasks [69, 71]. The circles and crosses represent the object and goal positions, respectively. Solid blue objects indicate training tasks, while empty brown objects indicate test tasks.
Moreover, even if an agent successfully models the inductive bias parametrically, it is improbable that the same parameterization will be reusable for qualitatively distinct test tasks.
In addressing the challenges of non-parametric task variability in meta-RL, our primary strategy involves decomposing each non-parametric task into a set of shared elementary subtasks. We then parameterize each task based on the types of subtasks that constitute it. Despite the non-parametric task variability, tasks may share elementary subtasks. For instance in Figure 1c, a "Pick-place" task can be decomposed into subtasks: "grip object" and "place object," while a "Sweep-into" task can be decomposed into subtasks: "grip object" and "push object." By employing the shared subtask parameterization, the policy can capitalize on the captured commonalities between non-parametric tasks to enhance training efficiency and generalization capabilities.
However, our approach of task parameterization based on a set of subtasks faces two primary challenges: the lack of prior information about (1) the set of elementary subtasks and (2) the decomposition of each task. To address these issues, we employ meta-learning for the subtask decomposition (SD) process using a Gaussian mixture variational autoencoder (GMVAE) [29, 8, 59]. Our GMVAE encodes the trajectory up to the current timestep into latent categorical and Gaussian contexts, which are trained to reconstruct the task's reward and transition dynamics [74]. We discover that the meta-learned latent categorical context effectively represents the subtask compositions of tasks under non-parametric variations. Consequently, the policy, using the learned subtask composition, can readily generalize to new tasks comprising previously encountered subtasks. To further enhance generalization to unseen compositions of familiar subtasks, we propose a virtual training (VT) process [35, 1] specifically designed for non-parametric task variability. We train the policy on imaginary tasks generated by the learned dynamics decoder, conditioned on hypothetical subtask compositions.
We evaluate our method on the Meta-World ML-10 and ML-45 benchmarks [71], widely used meta-RL benchmarks comprising diverse non-parametric robotic manipulation tasks. We empirically demonstrate that our method successfully meta-learns the shareable subtask decomposition. With the help of the subtask decomposition and virtual training, our method, without any offline demonstration or test-time gradient updates, achieves test success rates of  $33.4\%$  on ML-10 and  $31.2\%$  on ML-45, which improves the previous state-of-the-art by approximately 1.7 times and 1.3 times, respectively.
2.1 Meta-Reinforcement Learning
A Markov decision process (MDP),  $M = (\mathcal{S}, \mathcal{A}, R, T, T_0, \gamma, H)$ , is defined by a tuple comprising a set of states  $S$ , a set of actions  $\mathcal{A}$ , a reward function  $R(r_{t+1}|s_t, a_t, s_{t+1})$ , a transition function  $T(s_{t+1}|s_t, a_t)$ , an initial state distribution  $T_0(s_0)$ , a discount factor  $\gamma$ , and a horizon  $H$ .
The goal of meta-RL is to learn to adapt to a distribution of MDPs with varying reward and transition dynamics. At the start of each meta-training iteration, an MDP is sampled from the distribution  $p(\mathcal{M}_{\mathrm{train}})$  over a set of MDPs  $\mathcal{M}_{\mathrm{train}}$ . Each MDP  $M_{k} = (S, \mathcal{A}, R_{k}, T_{k}, T_{0,k}, \gamma, H)$  is defined with a unique reward function  $R_{k}$ , transition function  $T_{k}$ , and initial state distribution  $T_{0,k}$ . Unlike in a multi-task setup, the agent in the meta-RL setup does not have access to the task index  $k$  that determines the MDP dynamics. The training objective is to optimize the policy  $\pi_{\psi}$  with
parameters  $\psi$  to maximize the expected return across all MDPs:  $\max_{\psi} \mathbb{E}_{M_k \sim p(\mathcal{M}_{\mathrm{train}})} [\mathcal{J}_{\mathrm{pol}}(\psi)]$ , where  $\mathcal{J}_{\mathrm{pol}}(\psi) = \mathbb{E}_{T_{0,k}, T_k, \pi_{\psi}} \left[ \sum_{t=0}^{H-1} \gamma^t R_k(r_{t+1}|s_t, a_t, s_{t+1}) \right]$ . During meta-testing, standard indistribution meta-RL methods are evaluated on tasks sampled from the same distribution  $p(\mathcal{M}_{\mathrm{test}})$  as the training tasks, i.e.,  $\mathcal{M}_{\mathrm{train}} = \mathcal{M}_{\mathrm{test}} = \mathcal{M}$  (Figure 1a). In contrast, OOD meta-RL methods assume strictly disjoint training and test task sets, i.e.,  $\mathcal{M} = \mathcal{M}_{\mathrm{train}} \cup \mathcal{M}_{\mathrm{test}}$  and  $\mathcal{M}_{\mathrm{train}} \cap \mathcal{M}_{\mathrm{test}} = \emptyset$  (Figure 1b).
2.2 Bayes-adaptive Meta-Reinforcement Learning
Since the true task index  $k$  is not provided to the agent in the meta-RL problem setup, it is important to balance exploration and exploitation while learning about the initially unknown MDP. A Bayes-adaptive agent [38, 10, 16] achieves this balance by updating its belief  $b_{t}(R,T)$  about the MDP based on its experience  $\tau_{:t} = \{s_0,a_0,r_1,s_1,a_1,r_2,\dots ,s_{t - 1},a_{t - 1},r_t,s_t\}$ . The agent's belief over the MDP dynamics at time  $t$  can be represented as a posterior given the trajectory, i.e.,  $b_{t}(R,T) = p(R,T|\tau_{:t})$ . By augmenting the state with the belief to form a hyper-state space  $S^{+} = S\times \mathcal{B}$ , where  $\mathcal{B}$  is the set of belief, a Bayes-adaptive MDP (BAMDP) can be constructed. The objective of a BAMDP is to maximize the expected return within a meta-episode while learning, where a meta-episode consists of  $n_{\mathrm{roll}}$  rollout episodes (i.e.,  $H^{+} = n_{\mathrm{roll}}\times H$  steps) of the same MDP:
$$
\mathcal {J} _ {\mathrm {p o l}} ^ {+} (\psi) = \mathbb {E} _ {b _ {0}, T ^ {+}, \pi_ {\psi}} \left[ \sum_ {t = 0} ^ {H ^ {+} - 1} \gamma^ {t} R ^ {+} \left(r _ {t + 1} \mid s _ {t} ^ {+}, a _ {t}, s _ {t + 1} ^ {+}\right) \right], \tag {1}
$$
where  $T^{+}(s_{t + 1}^{+}|s_{t}^{+},a_{t},r_{t}) = \mathbb{E}_{b_{t}}\left[T(s_{t + 1}|,s_{t},a_{t})\right]\mathbb{I}(b_{t + 1} = p(R,T|\tau_{:t + 1}))$  is the transition dynamics and  $R^{+}(r_{t + 1}|s_t^+,a_t,s_{t + 1}^+) = \mathbb{E}_{b_{t + 1}}\left[R(r_{t + 1}|s_t,a_t,s_{t + 1})\right]$  is the reward dynamics of the BAMDP. The posterior belief update in the indicator function  $\mathbb{I}(\cdot)$  is intractable for all but simple environments.
VariBAD [74] solves the inference and posterior update of the belief by combining meta-RL and approximate variational inference [28]. At each timestep  $t$ , a recurrent encoder  $q_{\phi_h}$  encodes the experience  $\tau_{:t}$  into a hidden embedding  $h_t = q_{\phi_h}(\tau_{:t})$ . The approximate posterior belief over the dynamics can be represented as the parameters of a multivariate Gaussian distribution:  $b_t = (\mu_{\phi_z}(h_t), \sigma_{\phi_z}(h_t))$ , where  $\mu_{\phi_z}(\cdot)$  and  $\sigma_{\phi_z}(\cdot)$  are neural networks. The latent context  $z_t \sim \mathcal{N}(\mu_{\phi_z}(h_t), \sigma_{\phi_z}^2(h_t))$  is used to estimate the MDP dynamics:  $p_{\theta_R}(r_{j+1}|s_j, a_j, s_{j+1}, z_t)$  and  $p_{\theta_T}(s_{j+1}|s_j, a_j, z_t)$  for  $j = 0, \dots, H^+ - 1$ , including the past and future. Then the problem of computing the posterior over the dynamics  $p(R, T|\tau_{:t})$  reduces to inferring the posterior  $q_{\phi}(z_t|\tau_{:t})$ , where  $\phi = \{\phi_h, \phi_z\}$ . A separate policy network  $\pi_{\psi}(a_t|s_t, b_t)$  is trained to optimize the BAMDP objective in Eq. (1).
2.3 Non-parametric Task Variability
The term "non-parametric" in the context of task variability is introduced in the Meta-World paper [69]. It is used to distinguish the task variability in Meta-World manipulation tasks from more simplistic parametric variations exhibited in standard MuJoCo tasks, such as variations in target goal positions, directions, and velocities. However, the term "non-parametric" might lead to misunderstanding. Because it could suggest that there are no parameters available for task parameterization, even though we have effectively parameterized them in our approach. Our major breakthrough is in enabling this parameterization in terms of subtask compositions, which was difficult using earlier methods that relied on simple latent parametrization. Hence, while we retain the "non-parametric" terminology, we guide the readers to view the scope of our work through the lens of modularity or composition-based generalization [27].
In this section, we introduce our novel meta-RL method, named Subtask Decomposition and Virtual Training (SDVT), to handle non-parametric task variability. Our approach is based on decomposing each task into a set of elementary subtasks and parameterizing each task based on the composition of these subtasks. To achieve this, we use a Gaussian mixture variational autoencoder (GMVAE) to meta-learn the subtask decomposition process. In addition, we introduce a virtual training process that improves generalization to tasks with unseen compositions of subtasks.
(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  $(z)$  latent contexts. These contexts, which are trained to reconstruct the forward dynamics, are utilized by the policy network. This structure is also applied to the virtual training, as illustrated in Figure 3a, with an optional dispersion layer integrated. (b) An example of the learned task inference process within a meta-episode  $(H^{+} = 5000$  steps) on "Pick-place" is shown. We report the values for each dimension of contexts:  $\omega_{\phi_y}(h_t), \mu_{\phi_z}(h_t, y_t)$ , and  $\log \sigma_{\phi_z}(h_t, y_t)$ .
3.1 Subtask Decomposition (SD) with a Gaussian Mixture Variational Autoencoder
Our goal is to meta-learn a set of elementary subtasks and to meta-learn the decomposition of each task into a composition of these subtasks. The core of our method focuses on meta-learning the approximate subtask composition  $y_{t} \in \Delta^{K}$  sampled from a  $K$ -class categorical distribution, where  $\Delta^{K}$  denotes the  $K$ -dimensional probability simplex. For example with  $K = 3$ , we want the subtask composition  $y_{t}$  to be learned somewhat like (0.5, 0.5, 0.0) for "Pick-place" and (0.0, 0.5, 0.5) for "Sweep-into," where each dimension of  $y_{t}$  represents the weight corresponding to subtasks in the order of "place object," "grip object," and "push object." To capture such subtask information, we use a Gaussian mixture variational autoencoder (GMVAE) to represent the latent space of non-parametric tasks as a Gaussian mixture distribution. Each task is represented as a  $K$ -dimensional mixture proportion  $y_{t}$ , resulting in each class representing a unique subtask shared across different tasks. Thus, the distribution of  $K$  subtasks is modeled using a categorical distribution with  $K$  classes, where each class is associated with a Gaussian distribution. The learned subtask composition  $y_{t}$  represents the agent's belief at time  $t$  about the current task's decomposition into subtasks. It is crucial to note that we don't necessarily want  $y_{t}$  to be a one-hot embedding representing the subtask that the agent is solving at time  $t$ . We want  $y_{t}$  to represent a mixing proportion [56] over all possible subtasks that the agent believes to be relevant to the current task, including past and future as in Figures 2b and 4. This distinction is vital to the model's effectiveness in solving a range of tasks, as it allows for flexibility in subtask identification and generalization across different tasks.
Architecture Our full model in Figure 2a, which is based on the VAE of VariBAD [74], consists of three main components: the encoder, decoder, and policy networks parameterized by  $\phi = \{\phi_h,\phi_y,\phi_z\}$ ,  $\theta = \{\theta_z,\theta_R,\theta_T\}$ , and  $\psi$ , respectively.
(1) The encoder is defined as  $q_{\phi}(y_t,z_t|h_t) = q_{\phi_y}(y_t|h_t)q_{\phi_z}(z_t|h_t,y_t)$ . A recurrent network encodes the past trajectory  $\tau_{:t}$  into a hidden embedding  $h_t = q_{\phi_h}(\tau_{:t})$ . First, the categorical encoder  $q_{\phi_y}(y_t|h_t):\mathrm{Cat}(\omega_{\phi_y}(h_t))$  samples  $y_t$ , where  $\omega_{\phi_y}(h_t)\in \Delta^K$ . We use the Gumbel-Softmax trick [23] with a high temperature  $(\tau = 1)$  when sampling  $y_t$  to form a soft label. Then the multivariate Gaussian encoder  $q_{\phi_z}(z_t|h_t,y_t):\mathcal{N}(\mu_{\phi_z}(h_t,y_t),\sigma_{\phi_z}^2 (h_t,y_t))$  samples a continuous latent context  $z_{t}$ , which contains the parametric information of the subtasks, in addition to the categorical subtask information of  $y_{t}$ .
(2) The decoder is defined as  $p_{\theta}(\tau_{:H^{+}},y_{t},z_{t}) = p(y_{t})p_{\theta_{z}}(z_{t}|y_{t})p_{\theta_{R},\theta_{T}}(\tau_{:H^{+}}|z_{t})$ . It reconstructs the reward and transition dynamics for all transitions in the meta-episode  $\tau_{:H^{+}}$ , using the latent context  $z_{t}$  as in Eq. (4). We assume a uniform prior of subtask composition  $p(y_{t}): \mathrm{Uniform}(1 / K)$  and a Gaussian regularization  $p_{\theta_z}(z_t|y_t):\mathcal{N}(\mu_{\theta_z}(y_t),\sigma_{\theta_z}^2 (y_t))$ . This encoder-decoder architecture
allows both the approximate posterior  $q_{\phi}(y_t, z_t | h_t)$  and the prior  $p(z_t)$  to follow Gaussian mixture distributions.
(3) The policy network,  $\pi_{\psi}(a_t|s_t,b_t)$ , is trained separately conditioned on the belief  $b_{t} = (\mu_{\phi_{z}}(h_{t},y_{t}),\sigma_{\phi_{z}}(h_{t},y_{t}))$  that are parameters of the Gaussian context. In practice, we also provide the parameters of the categorical encoder  $\omega_{\phi_y}(h_t)$  to the policy, which we find to improve the performance. The parameters of the distributions  $(\omega_{\phi_y},\mu_{\phi_z},\sigma_{\phi_z},\mu_{\theta_z}$ , and  $\sigma_{\theta_z})$  are modeled as outputs of multilayer perceptrs (MLPs) as in Appendix C.4.
Objective We optimize the GMVAE to maximize the evidence lower bound (ELBO) for all time steps  $t = 0, \dots, H^{+}$  over the trajectory distribution  $d(M_{k},\tau_{:H^{+}})$  induced by the policy in MDP  $M_{k}$ :
$$
\operatorname {E L B O} _ {t} (\phi , \theta) = \mathbb {E} _ {d \left(M _ {k}, \tau ; H +\right)} \left[ \mathbb {E} _ {q _ {\phi} \left(y _ {t}, z _ {t} \mid h _ {t}\right)} \mathcal {J} _ {\mathrm {G M V A E}} \right], \tag {2}
$$
$$
\mathcal {J} _ {\mathrm {G M V A E}} = \alpha_ {R} \mathcal {J} _ {\mathrm {R - r e c}} + \alpha_ {T} \mathcal {J} _ {\mathrm {T - r e c}} + \alpha_ {g} \mathcal {J} _ {\mathrm {r e g}} + \alpha_ {c} \mathcal {J} _ {\mathrm {c a t}}. \tag {3}
$$
In addition to the reconstruction objectives  $\mathcal{I}_{\mathrm{R - rec}}$  and  $\mathcal{I}_{\mathrm{T - rec}}$ , we have additional regularization  $\mathcal{I}_{\mathrm{reg}}$  and categorical  $\mathcal{I}_{\mathrm{cat}}$  objectives:
$$
\mathcal {J} _ {\mathrm {R - r e c}} = \sum_ {j = 0} ^ {H ^ {+} - 1} \log p _ {\theta_ {R}} \left(r _ {j + 1} \mid s _ {j}, a _ {j}, s _ {j + 1}, z _ {t}\right), \quad \mathcal {J} _ {\mathrm {T - r e c}} = \sum_ {j = 0} ^ {H ^ {+} - 1} \log p _ {\theta_ {T}} \left(s _ {j + 1} \mid s _ {j}, a _ {j}, z _ {t}\right), \tag {4}
$$
$$
\mathcal {J} _ {\text {r e g}} = \log \frac {p _ {\theta_ {z}} \left(z _ {t} \mid y _ {t}\right)}{q _ {\phi_ {z}} \left(z _ {t} \mid h _ {t} , y _ {t}\right)}, \quad \mathcal {J} _ {\text {c a t}} = \log \frac {p \left(y _ {t}\right)}{q _ {\phi_ {y}} \left(y _ {t} \mid h _ {t}\right)}. \tag {5}
$$
The regularization objective  $\mathcal{I}_{\mathrm{reg}}$  minimizes the KL divergence between the learned posterior Gaussian distribution  $q_{\phi_z}(z_t|h_t,y_t)$  and learned Gaussian priors  $p_{\theta_z}(z_t|y_t)$ . Unlike the standard VAE that assumes a standard normal prior, we learn  $K$  distinct Gaussian priors conditioned on  $y_{t}$ . The categorical objective  $\mathcal{J}_{\mathrm{cat}}$  maximizes the conditional entropy of  $y_{t}$  given  $h_t$  and prevents collapse. Refer to Appendix B for the derivation of the ELBO objective. The reconstruction objectives in Eq. (4) are computed for all timesteps from the first to the terminal step of the meta-episode. Therefore, the subtask composition  $y_{t}$  at time  $t$  is not necessarily a one-hot label of the belief about the current subtask at time  $t$ , but a mixture label of the belief on all subtasks that compose the current task in the past and future within the meta-episode. Under the BAMDP setup, the agent learns to reduce the uncertainty of the decomposition as quickly as possible, supporting the policy with the converged parameterization. Combining the policy objective in Eq. (1) and the sum of the ELBO objectives in Eq. (2) for all timesteps in a meta-episode, the overall objective over training tasks is to maximize:
$$
\mathcal {J} (\phi , \theta , \psi) = \mathbb {E} _ {M _ {k} \sim p \left(\mathcal {M} _ {\text {t r a i n}}\right)} \left[ \mathcal {J} _ {\mathrm {p o l}} ^ {+} (\psi) + \sum_ {t = 0} ^ {H ^ {+}} \mathrm {E L B O} _ {t} (\phi , \theta) \right]. \tag {6}
$$
Occupancy regularization The dimension of  $y_{t}$  or the number of underlying subtasks  $K$  is a crucial hyperparameter that should be determined based on the number of training tasks  $N_{\mathrm{train}}$  and their similarities. However, in many cases, prior information about the optimal value of  $K$ , denoted as  $K^{*}$ , may not be available. One way to expand the scope of our method for unknown  $K^{*}$  is to meta-learn the number of effective subtasks as well. First, we assume  $K^{*} < N_{\mathrm{train}}$ , otherwise each task will be classified into a separate subtask with one-hot label, preventing learning shareable subtasks. We start with a sufficiently large  $K = N_{\mathrm{train}}$  and regularize the ELBO objective to progressively reduce the number of effective subtasks (non-zero components) occupied in  $y_{t}$  with the following occupancy regularization that penalizes the usage of larger indices in the subtask composition:
$$
\mathcal {J} _ {\text {o c c}} = - \log K (e ^ {- K + 1}, e ^ {- K + 2}, \dots , e ^ {- 1}, e ^ {0}) \cdot y _ {t}. \tag {7}
$$
We calculate the dot product of exponential weights and the subtask composition  $y_{t}$  to penalize the occupancy of higher dimensions of  $y_{t}$ . We scale the dot product by  $\log K$  to match the scale to the upper bound of  $\mathcal{I}_{\mathrm{cat}}$ . We add  $\mathcal{I}_{\mathrm{occ}}$  multiplied by a coefficient  $\alpha_{o}$  to the GMVAE objective in Eq. (3). Consequently, the agent prioritizes using lower indices in the decomposition to represent frequent subtasks and sparingly uses higher indices for rare subtasks as in Figure 4b and in Appendix E.1.
Decoder dropout As the GMVAE is optimized using the trajectories induced by the policy, the decoder can easily overfit the frequent states and actions of training tasks [35]. This can lead to low predicted rewards for unexperienced states and actions, regardless of the latent context  $z_{t}$ . When
(a) Imaginary reward generation.
(b)  $\tilde{y} = [1,0,0,0,0]$
(c)  $\tilde{y} = [0,1,0,0,0]$ .
(d)  $\tilde{y} = [\frac{1}{2},\frac{1}{2},0,0,0]$
(e)  $\tilde{y} = [0,0,\frac{1}{3},\frac{1}{3},\frac{1}{3}].$
Figure 3: Virtual training. (a) Generation of imaginary rewards with the decoder conditioned on a fixed imaginary subtask composition  $\tilde{y}$ . (b)-(e) Examples depicting the diversity of generated imaginary tasks, where  $K = 5$  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  $\tilde{y}$ .
such overfitting happens, the latent context loses its task-informative value, leading to the potential underperformance of policy based on this context. Following the approach of LDM [35], we address this by applying dropout (DO) to the state and action embeddings of the decoder while leaving latent context  $z_{t}$  untouched:  $p_{\theta_R}(r_{j + 1}|\mathrm{DO}[s_j,a_j,s_{j + 1}],z_t)$  and  $p_{\theta_T}(s_{j + 1}|\mathrm{DO}[s_j,a_j],z_t)$ . This dropout application is crucial for the virtual training discussed in the following section.
3.2 Virtual Training (VT) on Generated Tasks with Imaginary Subtask Compositions