Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks
Authors: Jeongmo Kim
Yisak Park
Minung Kim
Seungyul Han
Venue: ICML 2025 poster (2025)
Keywords: Reinforcement Learning
Deep Reinforcement Learning
Meta Reinforcement Learning
Representation Learning
Generalization
Machine Learning
Out-Of-Distribution Tasks
Primary area: reinforcement_learning->deep_rl
TL;DR: To improve generalization for out-of-distribution tasks in meta-RL, our Task-Aware Virtual Training algorithm proposes a metric-based task latents learning method, task-preserving learning method of virtual task, and a state regularization method.

Abstract: Meta reinforcement learning aims to develop policies that generalize to unseen tasks sampled from a task distribution. While context-based meta-RL methods improve task representation using task latents, they often struggle with out-of-distribution (OOD) tasks. To address this, we propose Task-Aware Virtual Training (TAVT), a novel algorithm that accurately captures task characteristics for both training and OOD scenarios using metric-based representation learning. Our method successfully preserves task characteristics in virtual tasks and employs a state regularization technique to mitigate overestimation errors in state-varying environments. Numerical results demonstrate that TAVT significantly enhances generalization to OOD tasks across various MuJoCo and MetaWorld environments. Our code is available at https://github.com/JM-Kim-94/tavt.git.

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks
Jeongmo Kim
Yisak Park
Minung Kim
Seungyul Han
Abstract
from off off-policy samples and uses these samples for rein- forcement learning, resulting in sample-efficient learning and faster convergence compared to traditional methods. In contrast, another prominent context- bbsed meta-RL method, VariBad (Zintgraf et al., 2019), employs Bayesian ap- proach to learn a belief distribution over environments, ef- fectively managing the exploration-exploiation trade-off in previously unseen environments. Recently, to address this issue and better distinguish tasks based on their char- acteristics, advanced representation learning methods such as contrastive learning have been increasingly adopted in meta-RL (Fu et al., 2021; Choshen & Tamar, 2023). CCM (Fu al., 2021) is a meta-RL method that integrates con- trastive learning with PEARL. In CCM, task latents from the same task treated having positive relationship and are trained to be close to each other, while latents from different tasks are treated as having a negative relationship and are trained to be distant.
Meta reinforcement learning aims to develop poli- cies that generalize to unseen tasks sampled from a task distribution. While context-based -based meta-RL RL methods improve task representation using task la- tents, they often struggle with out-o--distribution (OOD) tasks. To address this, we propose Task- Aware Virtual Training (TAVT), novel algorithm that accurately captures task characteristies for both training and OOD scenarios using metric- based representation learning. Our method suc- cessfully preserves task characteristies in virtual tasks and employs state regularization technique to mitigate overestimation errors in varying environments. Numerical results demonstrate that TAVT significantly enhances generalization to OOD tasks across various MuJoCo and Meta- World environments. Our code is available at http:/gghuuubcooo//MMKK-//-aaaaa
Advanced representation learning in meta-RI. improves the ability to distinguish training tasks through learned task latents. However, most existing methods assume that the test task distribution matches the training distribution, limit- ing their effectiveness on Latent Dynamics Mixture (LDM) (Lee & Chung, 2021) addresses this by training policies with virtual tasks (VTs) created via linear interpolation of task latents Icarned by VariBad, improving OOD generalization. Despite its bene- fits, we identify two key issues with the existing methods for VT construction. First, generated VTs often fail to capture task characteristics accurately. Second, LDM focuses solely on reward sample generation, which struggles in environ- ments with task-dependent state transitions. To overcome these limitations, we propose Aware Virtual Training (TAVT), novel algorithm that generates VTs accurately reflecting task characteristics for both training and OOD scenarios. Using Bisimulation metric (Ferns al., 2011; Ferns & Precup, 2014), our method captures task variations, such as changing goal positions, and incorporates an on-off task latent loss to stabilize the task latents. In addition, we introduce task peeserreenringg sample generation to ensure VTs generate realistic sample contexts while maintaining task- specific features. Finally, address state-varying environ- ments, task decoder generates full dynamics including
1. Introduction
(OOD) tasks.
Research in meta reinforcement learning (meta-RL) aims to train policies on training tasks sampled from a training task distribution, with the goal enabling the learned policy to adapt and perform well on unseen test tasks. Model- agnostic meta-learning (MAML) (Finn al., 2017) seeks to find initial parameters that can generalize well to new tasks by using policy gradients to measure how well the current policy parameter can adapt to new tasks. On the other hand, to effectively distinguish tasks while capturing the task char- acteristics, context-based meta-RL methods that learn task latents through representation learning has been researched recently (Rakelly et al., 2019; Zintgraf et al., 2019; Fu et al., 2021). One of the most well-known context- based methods, PEARL (Rakelly et al., 2019), learns task representations
'Graduate School of Artificial Intelligence, UNIST, UI- san, South Korea. Correspondence to: Scungyul Han <sy- han an
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Task-Aware Virtual Training: Enhancing Generalization in Learning for Out-of-Distribuui Tasks
PEARL
CCM
VariBad/LDM
TAVT (Ours)
(a)
(b)
(c)
(d)
(e)
Figure 1. (a) 2D Goal positions in the Ant-Goal environment: The blue shaded area indicates the training task distribution, with blue marks representing inner training tasks, green marks representing outer training tasks, and red marks denoting OOD test tasks. (b-c)
t-SNE visualization of task latents for various context- context-based meta-RL RL methods.
both rewards and next states, and we propose state reg- ularization method to mitigate overestimation errors from generated samples.
Po is the initial state distribution. At each time step the agent selects an action a based on the policy receives a reward rt == Rd (st, 0e). and transitions to the next state St+1 ~ p7 (-|s;,a;) The MDP for each task may vary, but tasks share the same state and action spaces. During meta- mneta-training, the policy is optimized to maximize the cumulative reward sum E, 2"rx across tasks sampled from p(Turain) The policy is then evaluated on OOD test tasks from p(Tust), which differs entirely from p( p(Terain).
Table compares our method with existing approaches, and Fig. illustrates differences in task latent representations for the Ant-Goal-OOD environment, where an ant agent should reach a goal point. Fig. ((a) shows 2D goal posi- tions, including training and OOD test tasks. While PEARL struggles to distinguish tasks, CCM, VariBad, and LDM differentiate training tasks but scatter OOD test task repre- sentations or fail to align goal position angles, In contrast, our method (TAVT) accurately aligns task latent for both training tasks and OOD test tasks, reflecting task character- istics more effectively. In particular, while other methods either disregard VT or rely on simple eeconstruction-basee VT. our approach utilizes metric-based representation and generative methods, enabling the encoder to learn more ac- curate VT contexts that effectively capture task information. This highlights the novelty of our proposed VT framework. To introduce the proposed TAVT, the paper outlines the meta-RL setup in Section 2, our approach in Section 4, and experimental results in Section 5, showing improved task representation and OOD generalization.
2.2. Context-based Meta RL
Recent meta-RL RR methods focus on learning latent contexts to differentiate tasks. PEARL (Rakelly et al., 2019), a well- known context-based meta-RL. approach, learns task latents z ~ ( using a task encoder 9s with parameters v and task context c := {(s), 01, r1, s where s' is the next state and N₂ is the number of transition samples. PEARL defines task-dependent policies <--|s, 2) and Q- functions ((s, a, z), using soft actor-critic (SAC) (Haarnoja et al., 2018) to train policies. The task encoder qy is trained to minimize the encoder loss:
Eeppaan [E [TQ
Table J. Comparison of Context-based Meta RL Methods
where N the multivariate Gaussian distribution, I the identity matrix, and Le is the SAC critic loss. Inspired by the variational auto-encoder (VAE) (Kingma & Welling, 2013), PEARL uses LQ instead of VAE's reconstruction loss to better distinguish tasks.
Task Representation
Virtual Tasks
Task peok-ineerving VT Sumples
PEARI.
O
x
X
VariBAD
CCM
0
X
LDM
(rowand only)
TAVT (QUi))
0
O
2.3. Virtual Task Construction
2. Preliminary
To improve the generalization of policies to various OOD tasks, LDM (Lee & Chung, 2021) defines the task latents za for a virtual task as a linear interpolation of training task latents z i = 1,... M,, expressed as:
2.1. Meta Reinforcement Learning
In meta-RI., each task T is sampled from task distribution P(T) and defined as a Markov Decision Process (MDP) (S, A, PT RT_" RT,77P0 Po). where 5 and A are the state and ac- tion spaces, PT represents state transition dynamics, RT is the reward function, ? E [0, 1) is the discount factor, and
(1)
2

---

Task-Aware Virtual Training: Enhancing Generalization in Meta- Reinforrrnfr Learning for Out-of- Tasks
where a (at..... . @Diricheet(1, 1,,1, 1,.... 1) 1-1 wf is the interpolation coefficient, Dirichlet(:) is the Dirichlet distribution, and M is the number of training tasks used for mixing. The parameter 5 2 1 controls the degree of mixing: with _ = 1, only interpolation within the train- ing task latents occurs, while 3 > allows extrapolation beyond the original latents. LDM further trains the policy using contexts generated from the task decoder based on the interpolated latents z", enabling it to handle OOD tasks more effectively.
4. Methodology
4.1. Metric-based Task Representation
In this section, we propose a novel representation learning method to ensure task latents accurately capture differences in task contexts, enabling virtual tasks to effectively re- flect task characteristics. To achieve this, we leverage the Bisimulation metric (Ferns et al., 2011), which measures the similarity of two states in an MDP based on the reward function Rd and state transition PT. In meta-learning, the Bisimulation metric can quantify task similarity by com- paring contexts (Zhang et al., 2021). Unlike Zhang et al. (2021), which considers tasks with different state spaces, we adapt the metric for tasks sharing the same state and action space, modifying it from Eq. (4) in Zhang et al. (2021).
3. Related Works
Advanced Task Representation Bearning:Advaneeee rep- resentation learning techniques have been widely explored to improve task latents that effectively distinguish tasks. Re- cent meta-RL methods use contrastive learning (Oord et al., 2018) to enhance task differentiation through positive and negative pairs, improving task representation (Laskin et al., 2020; Fu et al., 2021: Choshen & Tamar, 2023) and captur- ing task information in offline setups (Li al., 2020b; Gao et al., 2023). The Bisimulation metric (Ferns et al., 2011) is employed to capture behavioral similarities (Zhang et al., 2021; Agarwal et al., 2021; Liu et al., 2023) and group simi- lar tasks (Hansen-Estruch et al., 2022; Sodhani et al., 2022). Additionally, skill representation learning (Eysenbach et al., 2018) addresses non -onaammmttiicc meta-RL challenges (Frans
Definition 4.1 (Bisimulation metric for task representation). For two different tasks Ti and Ty.
d(Tr,T)
E
[RT (s, (R
(s.a)~
+ - pTs
(1) (1)
(2)
where D is the replay buffer that stores the sample contexts, RT. PT are the reward function and the transition dynamics for task T, W2 is 2- 2-Wasserstein distance between the two distributions, and v) E (0, 1] is the distance coefficient. Proposition 4.2. d(,,)) defined in Eq. (2) is a metric. Proof) The detailed proof is provided in Appendix A.
et al., 2017; Harrison et al., 2020; Nam et al., 2022; Fu et al., 2022; He et al., 2024), while task representation learning is increasingly applied in multi-task setups (Ishfaq et al., 2024; Cheng et al., 2022; Sodhani et al., 2021).
The Bisimulation metric d equals 0 when the contexts of two tasks perfectly match and increases as the context difference grows. Task latents learned using this metric more effectively capture task distances than existing rep- resentation methods, as shown in Fig. 1. We train the task encoder qu (z|c²) to ensure the task latent z ~ to preserves the Bisimulation metric d in the latent space. Since the actual reward function R and transition dynam- ies P are generally unknown, we train the task decoder (s a, z) = (Ra(s, a, z), (o a, z)) with parameters e to approximate task dynamics using reconstruction loss.
Generalization for OOD Tasks: Meta-RL techniques for improving policy generalization in OOD test environments have been actively studied (Lan et al., 2019; Fakoor et al., 2019; Mu et al., 2022). Model-based approaches (Lin et al., 2020; Lee & Chung, 2021), advanced representation learn- ing with Gaussian Mixture Models (Wang et al., 2023b; Lee et al., 2023), and Transformers ( (Vaswani et al., 2017; Melo, 2022; Xu et al., 2024) have been explored. Additionally, some studies tackle distributional shift challenges through robust learning (Mendonca et al., 2020; Mehta et al., 2020; Ajay et al., 2022: Greenberg et al.. 2023).
We adopt the learning structure of PEARL, which uses two distinct policies: Texp for on-policy exploration to obtain task latents from contexts and TRL for off- RL to maximize returns. Contexts generated by Texp and TRL are stored in the on-policy buffer D and the off-policy buffer D respectively. PEARL considers only on-policy task latents Zon derived from cT ~ DTT but zon can be unstable due to limited contexts in D To address this, we propose an on-off latent learning structure where off- policy task latents Zoff. derived from C ~ DT maintain the Bisimulation distance for training tasks, and Zon just aligns with Zofl- Fig. 2 shows t-SNE visualization of zon- demon- strating that on-off latent loss provides more stable task
Model-based Sample Relabeling: Model-based sample generation and relabeling techniques have gained attention in meta-RI. (Rimon et al., 2024; Wen et al., 2024), enabling the reuse of samples from other tasks using dynamics mod- els (Li et al., 2020a; Mendonca et al., 2020; Wan et al., 2021; Zou et al., 2024). These methods address sparse rewards (Packer al., 2021; Jiang et al., 2023), mitigate distribu- tional shifts in offline setups (Dorfman et al., 2021; Yuan & Lu. 2022; Zhou al., 2024; Guan et al., 2024), and in- corporate human preferences (Ren et al., 2022; Hejna III & Sadigh, 2023) or guided trajectory relabeling (Wang et al., 2023a), expanding their applications.
3

---

Task-Aware Virtual Training: Enhancing Generalization in Reinffeeinorreement Learning for Out-of-Diss cittof-Distribution Tasks
140
(a)
Tral OOD
Epoch
Direction
50
(a)
(b)
(b)
Figure 2. Changes in latents Zon of 4 randomly sampled training tasks in the Ant-Goal-OO seal-OOD environment: (a) Using Zon only (b)
Using on-off latent loss for task representation learning
representation compared to using Zon alone. In summary, our encoder-decoder loss is defined as
Figure 3. Comparison of task latents in the Ant-Dir-4 Environment: (a) Without task preserving loss (b) With task-preserving loss. The left graph shows the difference between latents x" and for VTs. during training. while the right graph shows latents for training (blue circles) and OOD tasks (red circles) during evaluation. Train- ing tasks have 4 goal directions of 5,, A = 0, 1, 2, 3, and OOD tasks cover 12 different directions from 0 2m.
Lhisim (v, 6) = BT.,TT pappppp(iii
[Raoet 2aal d(Tr. Ps:
Bisimulation loss
E (s.amaa) Di (t [(r F)² + (s' $1)²
Recoustinntion loss
(zin Zoa)
~ (-(" c. DTi, vi,
(3)
c8-off Intent loss
Despite the benefits of the task-presrrving loss, virtual con- texts ex may still differ from real contexts due to limitations in the task decoder (- z"), which cannot fully capture actual task contexts. These differences can introduce insta- bility and degrade performance for RL training. To address this issue, we use a Wasserstein generative adversarial net- work (WGAN) (Arjovsky al., 2017), designed to reduce the distribution gap between real and generated data. In our setup, the task decoder po(-, x²) acts as the generator, learning to produce samples that closely resemble real ones, while real contexts serve as the target for the discriminator. The discriminator fc increases its value for real samples and decreases it for generated samples, while the genera- tor aligns virtual contexts with real ones by increasing fc for generated samples. This process ensures that virtual contexts not only preserve task information but also closely resemble real contexts, significantly reducing the gap be- tween VT-generated samples and real OOD task samples. To summarize, the WGAN discriminator and generator losses, combined with the task-preserving loss, are defined as:
where d(Ti, Tj:Pz) replaces (RT,, PT.) in Eq. (2) with the task decoder pa( ): ) Vi, and 2 denotes gradient detach- ment for We construct the task latents 2 for VTs using the method in Section 2. Our proposed representation learn ing ensures task latents align more effectively to capture task differences based on the Bisimulation distance d, as illus- trated in 1ig.1. 1. Additionally, in Section 5.3 and Appendix B, we present task representations for other tasks, demon- strating that our method consistently aligns and stabilizes task latents across various environments.
4.2. Task Preserving Sample Generation
Using the task decoder pa(;22 2²) and VT task latents z we generate virtual contexts e := (s, fff 8/0 N. (())) 1=1
where (st, at) are sampled from real contexts c²,, and ,(" (s], at, z²). Existing VT construction methods (Lee & Chung, 2021; Lee al., 2023) use dropout-basee regularization to ensure virtual contexts generalize to unseen tasks. However, we observed that task latents 2" ~ (- (-(e)) often deviate significantly from zz, indicating that virtual contexts fail to effectively preserve task information. To address this, we propose a task- fask-preserving loss to minimize the difference between z" and 2², ensuring virtual contexts better retain task latent information. Fig. 3 demonstrates the effectiveness of the task- preeerreg loss. In Fig. 3(a), without the task-preserving loss, nd and 2 show significant differences, resulting inaccurate task latents for OOD tasks. In contrast, Fig. 3(b) shows that with the mask-preserving loss, z" and 2 align closely, enhancing stability and align- ment for OOD tasks.
Laise (C) =
Z(a)
(4)
+ Eea lC zo )1l + AGP Gradient Penalty, Egen (tb, () = Eeanpol I- l-sc 2
WGAN generator loss
E qu ( (119 (
(
Zoof ?11
(5)
task preserving
where the Gradient Penalty (GP), introduced in (Arjovsky et al., 2017), stabilizes training, with AGp as its coefficient. The detailed implementation of GP is provided in Appendix D.2. In this framework, the task decoder Ps is trained with

---

Task-Aware Virtual Training: Enhancing Generalization in Reinforeinfrr Learning for Out-of-Distrlbatoon Tasks
ples similar to real ones through task-preserving loss with WGAN. This enables the policy train diverse range of OOD tasks, improving generalization performance.
Tesp
bisim
Training Task:--.. Virtual Task:-
D
e
2
p
Since the agent cannot access off-policy latents for test tasks, it relies on on-policy latents zon obtained from Nexp episodes generated by the exploration policy Texps as in- troduced in PEARL (Rakelly et al., 2019). While PEARL uses a 2) with 2 (O. I) as the exploration policy, we propose novel exploration policy Texp ' (--1s, 2 that leverages VT task latents zon- This approach enables exploration across both training tasks and VTs by vary- ing the interpolation coefficient a, allowing the agent to explore a broader range of tasks. Appendix E shows that our exploration policy covers wider range of trajectories, including OOD tasks, leading to improved generalization. Fig. 4 provides an overview of the TAVT framework, with detailed implementation, including the RL loss function and meta-trainnigg/esting algorithms, in Appendix D.3.
TM
20
Duff
T~ (T)(T)
2
P
94
-gt
RT.
E yrn
ft
Edite
Figure 4. An illustration for the structure of TAVT
input Zoff. as described in Eq. (3). Both the task- pesk-preerrving loss and VT construction are always based on z derived from Zoff- Importantly, the off-policy task latent Zoft con- ditions sample context generation, with its gradient discon- nected to ensure stable training. By leveraging WGAN, we significantly reduce differences between generated and real samples, improving RL performance on OOD tasks. In Section 5.4 and Appendix C, we demonstrate that the pro- posed task- preserresng sample generation effectively bridges the gap between VT-generated samples and real OOD task
4.4. State Regularization Method
Gsu -10
loss VTs
Estimationn bies
200
2000
Cop 0:00
Cau 0.1
samples, enhancing generalization and RL performance.
4.3. Task-Aware Virtual Training
0
-100
0.4
0.6
0.8
1.0
0.4
0.6
0.8
1.0
By combining metric-based representation learning with task-preserving sample generation, we propose the Task- Aware Virtual Training (TAVT) algorithm, which enhances the generalization of meta-RI. to OOD tasks by leveraging VTs that accurately reflect task characteristics. The total encoder-decoder loss for TAVT is given by:
Time Steps
1e6
Time Steps
1e6
(a)
(b)
Figure 5. (a) Q-function loss for VTs (b) Estimation bias for OOD tasks in the Walker-Mass-OOD environment.
The virtual contexts ex in TAYT include both rewards and next states, enabling it to handle state-varying environments. However, inaccuracies in the task decoder can introduce errors in the Q-funncion,, leading to overestimation bias, a common issue in offline RL (Fujimoto al., 2019). While reward errors have minimal impact, next-state errors signif- icantly contribute to overestimation. To mitigate this, we propose a state regularization method, replacing the next states g in virtual contexts with 8/ a mix of next states s' from training tasks and g from virtual contexts:
Ltotal (V; Φ) = hisim (1b, b) + Egon( 1.,0 Φ),
(6)
with detailed loss scales provided in Appendix F.
Now, we train the RL policy TRL using SAC, leveraging both real contexts To D for each training task Ti and virtual contexts e generated by the proposed TAVT method. For training tasks, the RL policy is defined as TRI. [-1ss, z))) where z is the on-policy task latent. For VTs, the RL policy is defined as TRL. = T[- 2 where zo is the VT's task latent derived from - We train the RL policy TRI. using SAC, leveraging both real contexts Ti DT. for each training task T. and virtual contexts 8xx generated by the proposed TAVT method. For training tasks, the RL policy is defined as TAL 7 (--|s, z ). where i: is the on-policy task latent. For VTs,, the RL policy is defined as TRL = T( T(- ) where z is the VT's task latent derived from z The proposed TAVT method aligns task latents using metric-bace based representation and learns VTs that preserve task characteristics while generating sam-
ste = Ereg8"0 + - ereg)sss
where (8, a, (s,,,r,8') E cT, fa, ~ po(s, a, 2 and Greg E [0, 1] is the regularization coefficient.
Fig. 5(a) shows the Q-function loss for VTs during training, while Fig. 5(b) illustrates the Q-function estimation bias for OOD tasks in the Walker-Mass-OOD environment, where the agent's mass varies by task. Estimation bias, defined as the difference between Q-function values and average actual returns, significantly reduced with 6reg. = 0.1.
5

---

Task-Aware Virtual Training: Enhancing Generalization in Recta-eeiforeemment Learning for Out-of-Distribation Tasks
(a) Cheetah-Vel-OOD
(b) Ant-Dir-2
(c) Ant-Dir-4
(d) Ant-Goal-OOD
(e) Hopper-Mass-OOD
(f) Walker-Mass-O0
(g) ML1-Reach
(h) ML1-Push
Figure (a-f) MuloCo environments (g-h) MLI environments
In contrast, using Creg = 1.0 (full use of s*a) results in higher Q-function loss and greater bias. Section further demonstrates that this method improves OOD performance.
Table 2. Environmental setup
Musio
Must
Chectali Vel-OOD
10.0, 0.5) u (30,3.5)
(0.75, 1,25, 1,75, 2.2,,
5. Experiments
Ant-Dir-2
[0, 51
[4. 75)
Ant-Dir-4
[0, 1., T
(E+ in Br (r)
In this section, compare the proposed TAVT algorithm with various on-policy and off-policy meta-RL methods across MuJoCo (Todorov et al., 2012) and MetaWorid MLI (Yu et al., 2020) environments. For off-policy methods, we include PEARL (Rakelly et al., 2019), CCM (Fu et al., 2021) which uses contrastive learning to enhance represen- tation, MIER (Mendonca al., 2020) which incorporates brasee dynamics model, and Amago (Grigsby et al., 2024) which employs transformers for latent learning. For on on-policy methods, we evaluate MAML (Finn et al., 2017) which optimizes initial gradients for generalization, RL (Duan et al., 2016) which encodes task information in RNN hidden states, VariBad (Zintgraf et al., 2019) which applies Bayesian methods for representation learning, and LDM (Lee& & Chung, 2021) which uses virtual tasks with reconstruction loss and DropOut layer. Additionally, we provide an analysis of the task representation in TAVT as well as an ablation study on its performance. For TAVT, detailed hyperparameter configurations are provided in Ap- pendix F. For other methods, we reproduced their results using the author- author-provided source code and default hyperpa- rameters. Comparison graphs and tables show the average performance over random seeds, with standard deviations as shaded areas in the graphs or + in the tables.
Ant-Goal OOD
[0.0, L.0) u [2.5,
r= 1.75
8 E [0,2r]
9 E [0. 8.x. (0.75, 1.25, 1,75,
Cloper--Muss.Ce OOD
1.00 J (3.0,3.5)
Walke-Mas--OOD
10.0 U (30.
(0.75, 1.25, 1.75, 1,75,222,,275)
MI.I
M
M Minee
ML1-OOD
M
MLI-OOD-Ext
Muse
Mi
OOD, where only the reward function varies across tasks, and Salkeer-Mass-Oe and Hopper-Mass-OOD. where both the reward function and state transition dynamics vary. In addition, we considered 6 MetaWorla ML1 environments (Yu et al., 2020), including the original Reach and Push environments, as well as 4 OOD variations (Reach-OOD- Inter, Reach-OOD-Extra Push-OOD-Inter, and Push-OOD- Extra). these ML1 and MLI-OOD tasks, the final goal point varies across tasks while maintaining shared state dynamies. The task space is denoted by M, and detailed descriptions of these environments are provided below.
Cheetah-Vel-OODD The Cheetah agent is required to run at target velocities Star in M.
Ant-Dir: The Ant agent is required to move in direc- tions lair in M. For Ant-Dir-2, the training tasks in- volve 2 directions (fair = 0, 7 and for Ant-Dir-4, the training tasks involve 4 directions (fair = 0, 3. m, (T)
5.1. Environmental Setup
To evaluate generalization performance OOD tasks, we used 6 MuloCo environments (Todorov et al., 2012): Cheetah-Vel-OOD, Ant-Dir-2, Ant-Dir-4, and Ant-Goal-
Goal-OOD: The Ant agent should reach the 2D goal positions (rgoal cos 1 goal "goal sin 6goat), where the radius Tgoal and angle 8gaal are selected from M.
6

---

Task-Aware Virtual Training: Enhancing Generalization in Micfa-eeinforeemeet Learning for Out-of-Distribution Tasks
Cheetah-vel-oop
Ant- D
Ant-Dir-4
1200
1250
Test Avg Return
-100
1000
1000
800
750
-200
600
500
400
-300
250
200
8.0
0.5
1.0
1.5
2.0
Time Steps
1e6
Time Steps
1e6
Time Steps
1e6
Ant-Goal-OOD
Hopper-Mass-
Walker-Mass-OOD
600
800
-100
Test Avg Return
500
-200
400
-300
300
400
440
200
200
-500
100
0.2
0.4
0.6
0.8
10
0.0
0.2
0.4
0.6
0.8
1.0
Time Steps
Time Steps
Time Steps
TAVTIOurs) TAVT-oracle
VariBad
PEARL LDM
Amago MIER
MAML
Var-Badereeeee
RL
Figure 7. Performance comparison for MuloCo environments. The graphs for on-policy algorithms represent their final performance.
The Hopper/Walker agent are required to run forward with scale Mscall in M multiplied to their body mass.
shows the convergence performance of the test average return, while Table 1 presents the final average success rate for ML MLI/MLI-OOD environments. On-policy algo- rithms are trained for 200M timesteps in MuJoCo and 100M timesteps Meta environments. For off-policy al- gorithms, training timesteps vary in MuJoCo environments and are fixed at 10M timesteps for MetaWorld environments. Also, Fig. 7 includes the '-oracle' performance for VariBad (on-policy) and TAVT (off- loff-policy) in MuJoCo environments, This represents the upper performance bound achieved by training on all tasks from both Mtrain and Mtest.-
MLI/MLI-OOH The agent tasked with reaching or pushing an object a target goal position Itar, sampled from the 3D goal space M, which varies depending on the environment. To create OOD setup, inner areas Minner defined within the goal space M.
The training task space Merain and test task space Musst are detailed in Table 2. Except for the original MLI environ- ments, test tasks are OOD, meaning Mtrain MMtrett = 0. For Ant-Dir-2, test tasks include both interpolations be- tween training directions and extrapolations beyond them. In MI.1-0OD-Extra, OOD test tasks (iner) extrap- olate beyond training tasks (Minner). In all other OOD tasks, tasks lie in regions interpolating between training tasks. During training, tasks are uniformly sampled from Mtrainn with p(Tirain) = Unif(MMerin). A visualization of the enviromments is provided in Fig. 6. Further details on MuloCo and MetaWorld, including the definitions of the goal spaces M and Minner for MLI/MLI-OOD, are available in Appendix G.
For MuJoCo environments, Fig. shows that both VariBad and TAVT perform well in the '-oracle' setup across most en- vironments. However, in OOD task scenarios, TAVT signifi- cantly outperforms other on-policy and off-policy methods. Notably, TAVT achieves performance close to the "-oracle' in Cheetah- Cheetah-Vel-oee Vel-OOD, Ant-Goal-OOD, and Walker-Masss OOD, demonstrating strong generalization to unseen OOD tasks. While other algorithms struggle challenging envi- ronments like Ant-Dir-2, which includes extrapolation tasks, TAVT remains robust. In environments such as Walker- Mass- and Hopper-Mass-OOI. where state transition dynamics vary, LDM fails to adapt due to limitations in its VT construction. In contrast, TAVT excels, showcasing its ability to handle varying state transition dynamics. De- tails on policy traeectorie for training and OOD tasks are available Appendix E.
5.2. Performance Comparison
We compare the performance of the proposed TAVT with other meta-RL. algorithms. For MuloCo environments, Fig.

---

Task-Aware Virtual Training: Enhancing Generalization in Meii-oeiiforremment Learning for Out-of-Distributi Tasks
Table 3. Average success rate for MetaWorld MI.1 environments.
MAML
RL
VariBAD
LDM
PEARL
CCM
Amago
MIER
TAVT(Ours)
Reach
0.9740.02
0.95+0.04
0.73+0.12
0.76a0.1
0,48+0.21
0.65+0.13
0.71#0.27
0.61+0.18
0.98+0.02
Reach-OODDIInter
1.56+0.11
0.86+0.12
0.82+0.11 11
0.87+0.1
0.52+0.16
0.78+0.1
0.93=0.05
0.62+0.18
0.9620..03
Reach-OOD-Extra
0.48+0.15
0.73+0.14
0.82+0.11
0.79#0.15
0,48+0.14
0.81#0.12
0.43=0.08
0.65+0.12
0.99=0.01
Push
0.94+0.03
0.98+0.02
0.88=0.09
0.83+0.11
0.61
0.18+0.08
0.87=0.11
0.59+0.13
0.98=0.03
Pust-OOD-Inter
0.78+0.13
0.79=0.14
0.83+0.11
0.77=0.13
0.79=0.16
0.12#0.03
0.98=0.02
0,45#0.15
0.98=0.02
Push-OOD- COD--ODDEEEa Extra
0.55 13
0.
0.65+0.09
0.72+0.11 11
0.55+0.18
0.15+0.04
0.83=0.11
0.61
0.92+0.08
15
O n Tadks
Train Tacks
e 0000 ludkes
0:0 Taske
Mass
C fran Taons
Direction
GroRp sortion
C OCD Tapks
0.5
0.5
30
Tran Tasks OCD TRaKs
1D Goal Space
0
(a)
(b)
(c)
(d)
Figure 8. t-SNE visualization of task latents: (a) Cheetah-Vel-OOf (b) Walker-Mass-OOD (c) Ant-Dir-4 (d) Reach-OOD-Inter, For ML1
tasks in the 3D goal space, we provide both 3D representation for all tasks and 2D representation for tasks in the selected cross-scction.
For Meta World environments, Table 3 shows that TAVT consistently delivers superior performance across all setups. In the original ML1 environments, TAVT outperforms other methods, highlighting the advantages of its metric-based task representation. In MLI-OOD environments, PEARL- based methods such PEARL, CCM, and MIER perform poorly in Push and Push-OOD tasks, while other off-policy algorithms also fail to achieve success rates near 1. Despite being a PEARL-based method, TAVT achieves significantly higher success rates, often approaching 1. Compared to on-policy algorithms like LDM, RI.2, and VariBad, TAVT also demonstrates much better performance. These results highlight the effectiveness of our method in enhancing gener- alization for OOD tasks. The training graphs for MetaWorld environments are provided in Appendix H.1.
curately aligns task latents for both training and OOD test tasks, reflecting task characteristics. For instance, task la- tents align linearly by target velocity in and by target mass in Walker-Mass-ODD. Similarly, task latents align according to 2D target directions in Ant-Dir- 4 and 3D target goals in Reach-OOD-Iniee These results demonstrate that the proposed metric-based task representa- tion effectively distinguishes and generalizes task latents to OOD tasks. Task representations for other environments are provided in Appendix B. While t-SNE provides a reason- ably meaningful visualization for multi-dimensional latents, it does not preserve exact distances. To more accurately assess whether the learned latents preserve task-wise dis- tances under the proposed metric-based learning framework, we train TAVT with a 2D latent space and directly visual- ize the resulting representations. The results, presented in Appendix B.2, demonstrate that the 2D latents capture task characteristics comparably to the t-SNE projections while more faithfully preserving relative distances. This further validates the effectiveness of our metric-based representa- tion learning approach.
To enable more practical comparison, we report the com- putational cost of the proposed TAVT framework and the PEARL algorithm in Appendix H.2. While full TAVT requires approximately 18% more training time than the PEARL baseline due to the additional training of each com- ponent, other algorithms do not yield comparable perfor- mance gains even with similar computational overhead. This further demonstrates the practical advantage of TAVT.
5.4. Ablation Studies
Component Evaluation: To further analyze the proposed TAVT method, we perform component evaluation on the and Walker-Mass-OOW environments, comparing the final average return in Fig. 9. We consider the variations of TAVT including "TAVT w/o VT', which removes virtual tasks entirely. "TAVT w/o Egen', which ex- cludes the fgen component, "TAVT w/o on-off loss', which
5.3. Task Representation of TAVT
The comparative experiments confirm the superiority of TAVT. To explore the factors behind this improvement, Fig. 8 presents a t-SNE visualization of the task latents learned by TAVT across various environments. As shown in the cased of
environment in Fig. 1, TAVT ac-
8

---

Task-Aware Virtual Training: Enhancing Generalization in Micaa-Reinforeement Learning for Out-of-Distribation Tasks
Test Avg Return
-50
Test Avg Return
750
formance, reflected in Fig. underscoring the impor- tance of each component in TAVT. Notably, omitting the task-preserving sample generation loss Lgen significantly increases context differences. These findings demonstrate that the proposed WGAN-based task-preserving sample generation method effectively reduces context errors and enhances generalization performance, as detailed in Section 4.2. Additional context differences for other environments are provided in Appendix C.
-100
-150
-200
500
TAVT TAVT TAVTRecon
TAVT
Recon wro only on-off loss
wo
w/o only
Wo
wo
VT
Egen
on-off
VT
loss
(a)
(b)
Figure 9. Component evaluation on (a) Cheetah- Vel-OOD and (b)
State Regularization: the proposed state regularization method, Fig. 11 shows the performance across different Creg values. As illus- trated, a small Greg effectively reduces estimation bias. Fig. 11 further evaluates its impact on performance Walker- Mass-OOD and Hopper- SaapperMaas-OOOf environments, where Greg 0.1 yields the best results. This setting minimizes both Q-function loss and overestimation bias, enabling the Q-function to accurately reflect expected returns and im- proving performance in OOD tasks. In addition, we provide a finer sweep over Greg on the Walker-Mass- Mass-OOD environ- ment in Appendix C.2, where éreg = 0.1 still yields the best performance. These findings highlight the importance of the proposed state regularization method. Additional ablation studies the mixing coefficient for VT construction are provided in Appendix I.
To demonstrate the effectiveness of
Walker-Mass-OOD environments.
139
Context Difference
10-1
10a
10-2
10-3
10-1
0.4
0.6
0.8
1.0
0.4
0.6
0.8
1.0
Time Steps
les
Time Steps
leG
(a)
(b)
TAVT
TAVT w/o
TAVT w/o on-off loss
Recon only
Figure 10. Context differences between real contexts and virtual contexts generated the task decoder for OOD tasks: (a) Cheetah- Vel-OOD and (b) Malker-Mass-OOne environments.
600
6. Limitations
Test Avg Return
Test Avg Return
600
500
While TAVT shows strong performance through metric- based learning and task-aware sample generation, there areas for improvement. First, it involves several hyperpa- rameters that may require tuning for optimal performance, though it is not highly sensitive and our ablation studies provide practical guidance. Second, TAVT assumes a para- metric task distribution (e.g., velocity or mass changes in agent dynamics) and does not cover non-parametric distri- butions such as those ML10 or ML45 from Meta World, where task semanties differ entirely. This limitation could be addressed by combining TAVT with recent task decom- position methods (Lee et al., 2023), which we leave as promising direction for future work.
400
400
Kreg 0.0
Crep 0.5
200
300
Krep 1.0
Eng 0.1
200
0.0
0.2
0.4
0.6
0.8
1.0
0
Time Steps
1e7
Time Steps
1e6
(a)
(b)
Figure 11. Performance comparison for various éreg on (a) Walker-
Mass-OOD (b) Hopper-Mass Mass-OOD environments.
omits the on-off loss, and 'Recon only', which uses only reconstruction loss with DropOut in LDM. The results demonstrate that removing any component significantly de- grades performance, underscoring the importance of task representation learning and sample generation for OOD task generalization. In addition, TAVT outperforms the 'Recon only' approach, showcasing its ability to construct more effective virtual tasks than existing methods.
7. Conclusion
Existing meta-RL methods either overlook OOD tasks or struggle with them, particularly when state transitions vary. Our proposed method, TAVT, addresses these challenges through metric-based latent learning, task-preserving sam- ple generation, and state regularization. These compo- nents ensure precise task latents and enhanced generaliza- tion, overcoming previous limitations. Experimental results demonstrate that TAVT achieves superior alignment of OOD task latents with training task characteristics.
Context Differences: To evaluate the impact of context dif- ferences between VT-generated and actual contexts on per- formance, Fig. 10 presents these differences for the TAVT variations analyzed in the component evaluation. Since "TAVT w/o VT does not utilize virtual contexts, it is ex- cluded from Fig. 10. The results show that removing any component increases context error, leading to degraded per-

---

Task-Aware Virtual Training: Enhancing Generalization in Reii-feeiforreement Learning for Out-of-Distribation Tasks
Acknowledgment
ing via slow reinforcement learning. arXiv preprint arXiv: /61 1.02779, 2016.
This work was supported in part by the Institute of Informa- tion & Communications Technology Planning Evaluation (IITP) grant funded by the Korea government (MSIT), un- der the projects: No. 2022-0-00469 (Development of Core Technologies for Task-oriented Reinforcement Learning for Commercialization of Autonomous Drones), I1T- -TTP-2025-S-- 2022-0015636 (Innovative Human Resource Development for Local Intellectualization). and RS-2020-11200 01201000120136661111 (Arti- ficial Intelligence Graduate School Support UNIST); and in part by the 2021 Research Fund (No. 1.210149.01) of UNIST (Ulsan National Institute of Science & Technology).
Eysenbach, B., Gupta, A., Ibarz, J., and Levine, S. Diversity is all you need: Learning skills without reward function. arXiv preprint /002:v/800.060000 2018.
Fakoor, R., Chaudhari, P., Soatto, S., and Smola, A. J. Meta- q-learning, arXiv preprint
2019.
Ferns, N. and Precup, D. Bisimulation metrics are optimal value functions, In UAI, pp. 210-219, 2014.
Ferns, N., Panangaden, P., and Precup, D. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6):1622-1714, 2011.
Impact Statement
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. Interna- tional conference on machine learning, pp. 1126- 1135. PMLR, 2017.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here
Frans, K., Ho, J., Chen, X., Abbeel, P., and Schulman, J. Meta learning shared hierarchies. arXiv preprint arXiv: 1TXi::/10.097661 2017.
References
Agarwal, R., Machado, M. C.,, Castro, P. S., and Bellemare, M. G. Contrastive behavioral similarity embeddings for generalization in reinforcement learning, arXiv preprint arXiv:2101.05265, 2021.
Fu, H., Tang, H., Hao, J., Chen, C., Feng, X., Li, D., and Liu, W. Towards effective context for meta-reinforcement learning: an approach based on contrastive learning. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 7457-7465, 2021.
Ajay, A., Gupta, A., Ghosh, D., Levine, S., and Agrawal, P. Distributionally adaptive meta reinforcement learning- Advances in Neural Information Processing Systems, 35: 25856-25869, 2022.
Fu, H., Yu, S., Tiwari, S., Littman, M., and Konidaris, G. Meta-learning parameterized skills. arXiv preprint arXi::2206.0352 2022.
Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein gen- erative adversarial networks. International conference machine learning, pp. 223. PMLR, 2017.
Fujimoto, S., Meger, D., and Precup, Off-policy deep reinforcement learning without exploration. In Interna- tional conference machine learning, pp. 2052-2062. PMLR, 2019.
Brockman, G., Cheung, V., Pettersson, L., Schneider, J.,
Schulman, J., Tang, J., and Zaremba, W. Openai gym. arXiv preprint arXiv: 01500, 2016.
Gao, Y., Zhang, R., Guo, J., Wu, F., Yi, Q., Peng. S., Lan, S., Chen, R., Du, Z., Hu, X., et al. Context shift reduction for offline meta-reinforcement learning. Advances in Neural Information Processing Systems, 36, 2023.
Cheng, Y., Feng, S., Yang, J., Zhang, H., and Liang, Y. Provable benefit of multitask representation learning in reinforcement learning. Advances in Neural Information Processing Systems, 35:31741-31754 2022.
Greenberg, I., Mannor, S., Chechik, G., and Meirom, Train hard, fight easy: Robust meta reinforcement learn- ing. Advances in Neural Information Processing Systems, 36, 2023.
Choshen, E. and Tamar, A. Contrabar: Contrastive bayes- adaptive deep rl. In International Conference on Machine Learning, pp. 6005-6021 PMLR, 2023.
Grigsby, J., Fan, L.,, and Zhu, Y. Amago: Scalable in-context reinforcement learning for adaptive agents. In The Twelfth International Conference on Learning Representations, 2024.
Dorfman, R., Shenfeld, I., and Tamar, A. Offline meta reinforcement ceamiin--deentffabbiiyyy challenges and ef- fective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-4618, 2021.
Guan, C., Xuc, R., Zhang, Z., and Yu, Y. Cost-aware offline safe meta reinforcement learning with robust in disstiibutiionn online task adaptation. AAMAS, pp. 743-751, 2024.
Li, L., L., Li, Y.C., Yuan, L.,
Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I, and Abbeel, P. RIP: Fast reinforcement learn-
10

---

Task-Aware Virtual Training: Enhancing Generalization in Meta Reinforernoren Learning for Out-of-Distribation Tasks
Gulrajani, I., Ahmed, F,, Arjovsky, M., Dumoulin, V., and Courville, A. C. Improved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
decomposition. Advances in Neural Information Process-
ing Systems, 1443356-4338, 2023.
Li, J., Vuong, Q. Liu, S., Liu, M., Ciosek, K., Christensen, H., and Su, H. Multi-task batch reinforcement learning with metric learning. Advances in neural information processing systems, 33:6197-6211 2020a.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018.
Li, L., Yang, R., and Luo, D. Focal: Efficient fully- offline meta-reiiforcement learning via distance met- ric learning and behavior regularization. arXiv preprint arXiv:20 12, 2020b.
Hatsen-Estruch, P., Zhang, A., Nair, A., Yin, P., and Levine, S. Bisimulation makes analogies in goal-conditioned reinforcement learning. In International Conference Machine Learning, pp. 8407-8426. PMLR, 2022.
Lin, Z., Thomas, G., Yang. G., and Ma. T. Model-based ad- versarial meta-reinforcement learning. Advances in Neu- ral Information Processing Systems, 33:10161-10013, 2020.
Harrison, J.. Sharma, A., Finn, C., and Pavone, M. Con- tinuous meta-learning without tasks. Advances neural 2020. He, H., Zhu, A., Liang, S., Chen, F., and Shao, J. De- coupling meta-reinforcement learning with gaussian task contexts and skills. In Proceedings of the AAAI Confer- ence on Artificial Intelligence, pp. 2335-122366 2024.
Liu, Q., Zhou, Q. Yang, R., and Wang. J. Robust represen- tation learning by clustering with bisimulation metrics for visual reinforcement learning with distractions. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 8843-8851, 2023.
information processing systems, 33: 33:17571-17555
Mehta, B., Deleu, T., Raparthy, S. C., Pal, C. and Paull, L. Curriculum in gradient- basedd meta-reinforcemenn learn- ing. arXiv preprint ar 2ix::2002.00956. 2020.
Hejna III. D. J. and Sadigh, D. Few-shot preference learn- ing for numan-in-thee Learning, pp. 2014-2025. PMLR, 2023.
rl. In Conference on Robot
Melo, L. C. Transformers are rentrreiiforremmnnt learn- ers. In international conference on machine learning, pp- 15340-15359, PMLR, 2022.
Ishfaq, H., Nguyen-Tang, T., Feng, S., Arora, R., Wang, M., Yin, M., and Precup, D. Offline multitask representa- tion learning for reinforcement learning. arXiv preprint arXiv:2403. 11574, 2024.
Mendonea, R., Geng, X., Finn, C.,, and Levine, S. Meta- reinforcement learning robust to distributional shift via model identification and experience relabeling, arXiv preprint arXiv:2006.07177 2020.
Jangg Y., Kan, N., Li, C., Dai, W., Zou, J., and Xiong, H. Doubly robust augmented transfer for meta-reinnorcemment learning. Advances in Neural Information Processing Systems, 36, 2023.
Mu, Y., Zhuang, Y., Ni, F., Wang. B., Chen, J., Hao, J., and Decomposed mutual information opti- mization for generalized context renforeiinorremmnnt learning. Advances in Neural Information Processing Systems, 11.27563-27775 2022.
Luo, P. Domino:
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint a3Xiv:1312.6114, 2013.
Nam, T., Sun, S.-.., Pertsch, K., Hwang, S. J., and Lim, J. Skill-based reifrreiiforeement learning. arXiv preprint arXiv:2204.11822 2022.
Lan, L., Li, Z., Guan, X., and Wang, P. Meta reinforcement learning with task embedding and shared policy. arXiv preprint arXiv: 00199995506665 2019.
Oord, A. v. d., Li, Y., and Vinyals, O. Representation learn ing with contrastive predictive coding. arXiv preprint arXiv: A.Kiv/807.000 /807.03748, 2018
Laskin, M., Srinivas, A., and Abbeel, P. Curl: Contrastive unsupervised representations for reinforcement learning. In International conference muchine learning, pp. 5639-5650. PMLR, 2020.
Packer, C., Abbeel, P., and Gonzalez, E. Hindsight task relabelling: Experience replay for sparse reward meta-rl. Advances in neural information processing systems, 34: 2466 2477, 2021.
Lee, S. and Chung, S.-Y. Improving generalization in meta- with imaginary tasks from latent dynamics mixture. Advances in Neural Information Processing Systems, 34: 27222-27235, 2021.
Rakelly, K., Zhou, A., Finn, C., Levine, S., and Quillen, D. Efficient off-policy meta-reinforceement learning via prob- abilistic context variables. In International conference on machine learning, pp. 5331-5340 PMLR, 2019.
Lee, S., Cho, M., and Sung, Y. Parameterizing non- parametric reiforreito learning tasks via subtask

---

Task-Aware Virtual Training: Enhancing Generalization in Meta- ReiefoReiforremment Learning for Out-of-Distribation Tasks
Ren, Z.., Liu, A., Liang, Y., Peng, J., and Ma, J. Efficient meta reinforcement learning for preference-based fast adaptation. Advances in Neural Information Processing Systems, 35:15502-15515, 2022.
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. Meta-world: A benchmark and evalua- tion for multi-task and meta reinforcement learning, In Conference on robot learning. pp. 1094-1100 PMLR, 2020.
Rimon, Z., Jurgenson, T., Krupnik, O., Adler, G., and Tamar, A. Mamba: an effective world model approach for meta- reinforcement learning. arXiv preprint arXiv:2403.09855 2024.
Yuan, H. and Lu, Z. Robust task representations for of- fline meta-reinforcement learning via contrastive learning. In International Conference on Machine Learning, Pp. 25747-25759. PMLR, 2022.
Sodhani, S., Zhang, A., and Pineau, J. Multi-task rein- forcement learning with context- representations. In International Conference on Machine Learning, pp. 9767-9779. PMLR, 2021.
Zhang, A., McAllister, R. T., Calandra, R., Gal, Y., and Levine, S. Learning invariant representations for rein- forcement learning without reconstruction. In Interna- tional Conference on Learning Representations, 2021.
Sodhani, S., Meier, F., Pineau, J., and Zhang, A. Block contextual mdps for continual learning. In Learning for Dynamics and Control Conference, pp. 608-623, PMLR, 2022.
Zhou, R., Gao, C.-X., Zhang, Z., and Yu, Y. Gener- alizable task representation learning for offline meta- reinforcement learning with data limitations. In Proceed- ings of the AAAI Conference on Artificial Intelligence, pp. 17132 17400, 2024.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ inter- national conference on intelligent robots and systems, pp- 5026-5033. IEEE, 2012.
Zintgraf. L., Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hof- mann, K., and Whiteson, S. Varibad: A very good method for
deep rl via meta-learning. arXiv preprint 1rxiv: 2019.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention all you need. Advances in neural information processing systems, 30, 2017.
Zou, Q., Zhao, X., Gao, B., Chen, S., Liu, Z., and Zhang, Z. Relabeling and policy distillation of hierarchical re- inforcement learning. International Journal of Machine Learning and Cybernetics, pp. 1-17, 2024.
Wan, M., Peng, J., and Gangwani, T. Hindsight foresight relabeling for meta-reinnorcemment learning. In Interna- tional Conference on Learning Representations, 2021.
Wang, L., Zhang, Y., Zhu, D., Coleman, S., and Kerr, D. Supervised meta rentrreiiforre learning with trajectory optimization for manipulation tasks. IEEE Transactions on Cognitive and Developmental Systems, 16(2):681-691, 2023a.
Wang, M., Bing, Z., Yao, X., Wang, S., Kai, H., Su, H., Yang, C., and Knoll, A. Meta-reinforcemeet learning based on self-supervised task representation learning. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 10157-10165, 2023b.
Wen, L., Tseng, E. H., Peng, H., and Zhang, S. Dream to adapt: Meta reinforcement learning by latent context imagination and mdp imagination. IEEE Robotics and Automation Letters, 9(11):901--9788, 2024. doi: 10.
(A99LRRA.022.3417111
Xu, T., Li, Z., and Ren, Q. Meta-reinforcement learning robust to distributional shift via performing lifelong in- context learning. In International Conference on Machine Learning, 2024.
12

---

Task-Aware Virtual Training: Enhancing Generalization in Meda-Reiiforremment Learning for Out-of-Distributii Tasks
A. Proof
Definition 4.1 (Bisimulation metric for task representation), For two different tasks T; and T
d(Ti: T₁
(s, - (s, a)l + (W2(P" ('1ss,,)))
(1)
(s,a)~
where D is the replay buffer that stores the sample contexts, R² PPT are the reward function and the transition dynamics for task T, W2 is 2- 2- distance between the two distributions, and E (0, 1] is the distance coefficient.
Proposition 4.2. d(: defined in Eq. (2) is metric.
Proof of Proposition 4.2
As mentioned in Section 2, each task T is represented by an MDP AT RT P7)7.))). the action space, PT represents the state transition dynamics, RD is the reward function, E [0, 1) is the discount factor, and po is the imitial state distribution. Since all tasks shares the same state space and action space in this paper, so T; T, and only if PT: (s'|s ss, a)) TT (s'|s, s, a) and RT a) (s, a), Vs, s' E S, a E A for any tasks Te, Ta E M. Thus, from the definition of d given by Eq. (2), d metric since d satisfies the following axioms:
where S is the state space. A is
1. (Non-negativity) d(T;, TT)) 2 0 since
|nd Ws( :) are non-negative,
2. d(Ti, = + T R and pT: PT) = T
3. (Symmetry) d(T:. T,) = =d(T,, Ti) from the definition.
4. (Triangle Inequality) d(T;; Tk) d(T,, T,) + d(TT;TT)) Tx):
d(Ti, Tk) E (
1 (s.) (")))a))))((())))
(s, RT (x, (17 as,
(W2((Pi(11s,0)) (1),,a)) ( p.h (m)
EOODDD R [
RT (8, (1. (T5
B(Sa)DD [RT RT (s. ( (1)s p (7)
d(T,, T₁) + d(T,, Tk);
where can be derived, as both the absolute value | and the Wasserstein distance W2 satisfy the triangle inequality.
13

---

Task-Aware Virtual Training: Enhancing Generalization in Mede-Reinforrement Learning for Out-of-Distributis Tasks
B. Visualization of Task Latents Representation
B.1. Visualization of Task Latents Representation Across All Environments
To present the task latent representation results for considered OOD environments, Fig. B.1 shows the task representations for 6 MuJoCo OOD environments CChetaa-Vel-OOOD, Ant-Dir-2, Ant-Dir-4, Hopper-Mass- OOD, and Galker-Maas-Onne and Ant-Goal-OOD), and Fig. B.2 illustrates the task representations for 4 MLI-OOD environments (Reach-OOD-Inter, Reach-OOID-Extnn Push-OOD-Interr, and Push-OOID-Extrra) All task representations in Fig. B.1 and Fig. B.2 obtained during the meta-testing phase. Additionally, to achieve stable task representations for Fig. B.2, larger number of training tasks (Nerain = 500) was sampled compared to the originally required number (Nirain 50).
From the results in Fig. B, 1 the task latent representations for Cheetah- Vel-hO,, Ant-Dir-2, Ant-Dir-4, Hopper-Mass-OOD, and Sarkker-Mass-OOee are well-aligned according to their respective task characteristics, such as target velocity, target directions, and agent mass. For Ant- Ant-Goal- Goal-OOD. the task representation is aligned based on two characteristics: goal radius and direction. These findings demonstrate that the proposed method effectively aligns task representations with task characteristics across all MuJoCo environments, while OOD tasks maintain latents that preserve these characteristics.
Similarly, the results in Fig. B.2 show that for each MLI-OOD environment, the 3D visualization of the target goal positions is presented on the left, and the task representations aligned with the goal positions are displayed on the right. The results indicate that for all considered environments, both training and OOD tasks are well-aligned in 3D space according to their target goal positions. These findings highlight that the proposed Bisimulation metrie-based task representation learning effectively creates task latent space aligned with task characteristics, while the TAVT training method enhances the generalization of OOD task representations.
Train Tasks Tasks
Train Tasks
Velocity
Direction
Direction
OOD Tasks
0.0
(a) Cheetah-Vel-OOD
(b) Ant-Dir-2
(c) Ant-Dir-4
Direction
2.1
3.0
3.5
3.5
Train Tasks
Tasks
Train Tasks
0O0 Tasks
000 Tasks
25
3.
Radius
Mass
Mass
20
0.0
0.00
0.0
(d) Ant-Goal-OOD
(e) Hopper-Mass-OOD
(f) Walker-Mass-OOD
Figure B.I. (a-f) t-SNE visualization of the 6 MuloCo OOD environments
14

---

Task-Aware Virtual Training: Enhancing Generalization in Redt-Reeiforrementt Learning for Out-of Dut-of-Distribation Tasks
(a) Goal space (OOD-Inter setup)
(b) Reach-OOD-Inter
(c) Push-OOD-Inter
(d) Goal space (OOD-Extra setup)
(e) Reach-O00 -0t--raaa
(f) Push-OOD-Extra
Train Tasks
00D Tasks
Low
Goal High
Figure B.2. (a, d) Visualization of the goal space of each MLI-OOD setup (b, d. f) t-SNE visualization of the ML1-OOD environments
B.2. Visualization of 2-Dimensional Task Latents Representation
While t-SNE visualization method provides a reasonably meaningful visualization for multi-dimensione latents through effective manifold learning algorithm, it does not preserve exact distances between the latent vectors. To more accurately assess whether the learned latent space actually preserves task-wise distances under the proposed metric-based learning framework, we train TAVT with a 2D latent space on Ant-Gout Goal-ooDOOOD enviromment and directly visualize the resulting latent representations on 2D plane, enabling to reflect the direct distance between latent vectors. Fig. B.3(a) shows the Ant-Goal-OOD environment setup that is introduced in Fig. ((a), and Fig. B.3(b) shows the direct visualization of 2-dimensional task latents. Fach axis represents each element 2D latent vector. Only the latent dimension changes from 10 to from Fig. I(c) setup. This direct visualization indicates the task latent space tend to generally reflect the task geometry and preserve the task-wise distance despite very low dimensional latent vector including both training tasks and OOD test tasks. This demonstrates that the 2D latents capture task characteristics comparably the t-SNE projections while faithfully preserving relative distances and validates the effectiveness of our metric-based representation learning approach.
2D-Latent Representation on Ant-Goal-OOD
Second Element of Leanred Latent Vector
0.5
0.0
-05
-1.5
-1.0
-0.5
0.0
1.0
Element
Vector
(a)
(b)
Figure B.3. Visualization of 2-dimensional task latent on Ant-Goal-OOD environment: (a) 2D Goal positions the Ant-Goal-OOD
environment, which is introduced in Fig. (()). (b) Direct visualization of 2D latent vectors 2D plane.
15

---

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcemmnt Learning for Out-of-Distributie Tasks
C. Ablation Studies of the Task-Preserving Sample Generation.
C.1. Effectiveness of the Proposed Task-Preserving Sample Generation
In Section 4.2, we introduce a preeerrrsgg sample generation technique to ensure that virtual contexts retain task latent information while closely resembling actual task samples. To evaluate the effectiveness of this approach, we analyze how virtual sample contexts generated by the task decoder differ from real task contexts for OOD tasks. Context difference graphs for all considered OOD setups provided, showing the average reward and state differences between real and generated contexts for both MuJoCo and ML1 environments.
Fig. C.I presents context differences across 6 MuJoCo enviromments under OOD setups. Our TAVT algorithm achieves the smallest context differences in all environments, with the differences being most pronounced Cheetaa-Vell-OOl Ant-Goal-OOD, and Ant-Dir-4. Similarly, Fig. C.2 shows context differences in 4 MLI-OOD environments, where TAVT again demonstrates the smallest differences. These results highlight TAVT's effectiveness generating accurate transition samples, even for unseen OOD tasks.
In comparison, the 'Recon only' and "TAVT w/o on-off loss' setups exhibit the largest context differences, followed by "TAVT w/o Egen', while the proposed TAVT method achieves the smallest context difference. The 'Recon only' method suffers due to its reliance on transition samples solely from training tasks. The absence of the proposed on-off loss or Zgen leads to increased context differences, demonstrating that these components significantly reduce errors in the sample contexts generated by the proposed VT. This reduction context errors directly contributes to improving OOD task generalization.
10a
Context Difference
10-1
10a
15
10
10-1
10-3
0.4
0.6
0.8
1.0
2.5
3.0
3.5
4.0
4.5
5.0
0.4
0.6
0.8
1.0
Time Steps
1e6
Time Steps
1e5
Time Steps
1e6
(a) Cheetah-Vel-OOD
(b) Ant-Dir-2
(c) Ant-Dir-4
Context Difference
10-1
10a
N
10-2
10
LO-1
0.6
0.4
0.6
0.8
1.0
0.6
Time Steps
1e6
Time Steps
1e6
Time Steps
1e6
(d) Ant-Goal-OOD
(e) Hopper-Mass-OOD
(f) Walker-Mass-OOD
TAVT
TAVT w/o Egen
TAVT w/o on-off loss
Recon only
Figure C. (a-I) Context differences between real contexts virtual contexts generated by the task decoder the six MuloCo OOD tasks. The y-axis is plotted on log scale except Ant-Dir-2 environment to show the difference.
16

---

Task-Aware Virtual Training: Enhancing Generalization in Reinforrrnor Learning for Out-of-Distribation Tasks
Context Difference
10a
10
100
10
10-1
10-1
10a
10
10-2
10a
10-3
10-2
10
10
10-1
0.8
1.0
0.4
0.6
0.8
0.5
1.0
Time Steps (a) ML1-Reach-OOD-inter
Time Steps
le6
Time Steps
le6
Time Steps
(b) ML1-Reach-OOD--xxrr
(c) ML1-Push-OOD-Inter
(d) ML1-Push-OOD-Extra
TAVT
TAVT w/o Cpen
TAVT w/o on-off loss
Recon only
Figure C.2. (a-d) Context differences between real contexts and virtual contexts generated by the task decoder for the four MLI OOD
tasks. The y-axis is plotted on log scale show the difference.
C.2. Additional Ablation study of éreg
Section 4.4, we propose state regularization method to mitigate the overestimation problem that occurs when learning the Q-function using generated virtual states. In this section, we extend our ablation study include finer sweep over Ereg [0.0, 0.05, 0.1, 0.2, 0.5, 1.0] on the Walker-Mass-OOD environment. Fig. C.3 shows the additional ablation experiment results for Creg in the Walker-Mass-OOn environment. Fig. C.3(a) presents the performance according to Greg* exhibiting concave trend where relatively high performance is achieved at smaller values of frg,, with the best performance occurring at freg = 0.1. In addition, Fig. C.3(b) shows the estimation bias of the Q-function according to Ereg. defined as the difference between Q Q-function values and average actual returns. When Crgg = 0.05 or Creg = 0.1, the estimation bias of the Q- Q-function the smallest, while the others show a greater degree of overestimation or underestimation. These results confirm that small values around 'rgg 0.1 consistently lead to lower Q-function estimation bias and improved return and imply that the proposed state regularization method effectively helps the Q-function learn from the generated virtual states.
300
Test Avg Return
Estimation blas
200
0.0
500
Ereg 0.05
400
100
Ereg 0.1
Ereg 0.2
300
Ereg 0.5
0
200
Ereg 1.0
0.2
0.4
0.6
0.8
1.0
0.4
0.6
0.8
1.0
Time Steps
1e7
Time Steps
1e6
(b)
Figure C.3. Additional ablation study of erog on Walker Massker-Mas-OOOD environment: (a) Performance comparison (b) Estimation bias for
OOD tasks for various Ereg.
17

---

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcemeee Learning for Out-of- Tasks
D. More Detailed Implementations
In this section, we provide more detailed implementations for the proposed TAVT. Section D.1 describes the practical implementation of the loss terms in Chisim. Section D.2 explains the implementation of the gradient penalty for WGAN loss, and Section D.3 details the implementation of meeta-RL with TAVT, including the meta-training and meta-testing algorithms.
D.I. Practical Implementation of Liisim
To improve task representation and reflect task differences, we use the Bisimulation metric as described in Section 4.1. We also introduce an on-off loss to stabilize task latents. The encoder-decoder loss is given by Eq. (3), and to compute d(T, Ty; Pá) in Eq. (3),
the task decoder ( 2oft). However, since Zoft evolves during encoder training, the task decoder may become unstable, potentially affecting the metric d. To address this instability, we propose using additional decoder P5(: (-, idxx; with parameter e to compute the metric d, where idxy, is the one-hot encoded vector of the task index i. The decoder P3 is also trained using reconstruction loss. Consequently, the updated encoder-decoder loss is given by
Lisii Chisim(1,, ())) (TT.Typp((Tm))
2 - d(T,, Tr:pa)) + B (a ( (r (r- + (s' 8)) (Iraa
Bisimulaaiio
Reconstruction loss for
+ E (a,ar,"")) (1)pppollo
f)2 + (s' 81) + (zin - zoa) )
z ~ (-(c)) J. ~ DT. Vi. (D.I)
on-off latent loss
Reconstruction loss p
Moreover, to further stabilize the learning of the latent variables, instead of using single sample of 2 in the on-off loss, we use the average of 2o from multiple contexts sampled from the buffer. This approach helps prevent fluctuations in Z due to varying contexts, thereby aiding in more stable learning of the task latents.
D.2. Implementation of Gradient Penalty in Ldisc.
To stabilize the training of the adversarial network, the improved WGAN framework (Gulrajani et al., 2017) incorporates gradient penalty (GP) that restricts the gradient to prevent the discriminator from learning too quickly, as described in Section 4.2. GP ensures that the discriminator satisfies the I-Lipschitz continuity condition. In WGAN discriminator loss Eq. (4), the Gradient Penalty term is calculated by
Gradient Penalty = ET (erson) [(I| binter f (finter)) 1)
where dinter := - .z, + () zore) represents the interpolated samples between (c 3. which are induced by training tasks Tú, and (6 2. which are induced by the task decoder of VT. Here, 8 Unif([O, 1) is the interpolation factor for the GP. In addition, the WGAN structure trains the discriminator more frequently than the generator at a 5:1 ratio, as proposed in the original WGAN paper (Gulrajani al., 2017).
D.3. Meta-RL with TAVT and TAVT Algorithms
As described in Section 4.3, we aim to perform meta-RL using the contexts c/ obtained from training task Te and the virtual contexts 8 generated by the task decoder of the VT. For meta RL, the Q-function and the RL policy are defined as Qo(s, z) and Te(- z) with task latent z, where 0 represents the parameters of both the Q-function and the policy Based on SAC (Haarnoja et al., 201 8), the RL. losses for the Q-function and the policy are then given by:
Lo(0; c, z) = Bi (Qo(s, a, - (r- >rew + +.Ea ~ms 100 (s', ()) dent log (-1s, 2)1)
(D.2)
exp(Qo(s, ; 2 Aen))
c0; = Eane DKI. TT(-115522 2)
(D.3)
where Qo is the target value network with parameters 0- updated from 0 using the exponential moving average (EMA), Arew is the reward scale, Aent is the entropy coefficient, DKr. (pllq) is the Kullback-Leibler divergence between two distributions p and q, and Zg is the normalizing factor. Based on the proposed TAVT, we update the f-fnctioon and the RL policy with training task latents z using off-policy contexts obtained from these tasks, and update the Q- Quunciio and the
18

---

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforreemnn Learning for Out-of-Distribation Tasks
RL policy with the latents z of VT using virtual contexts generated by the task decoder, explained in Section 4.3, The total Q-function loss and policy loss with TAVT are given by
(0) E TT R go L D Lo(0 zon c AVT 20 Cart
(D.4)
Cx = E To-p(Frat) go. D C (0: zou coft) + AvT (0: 2 Ean)]
(D.5)
where ÄVT E [0, 1] is the loss coefficient for VT training. If XvT = 0, then TAVT does not utilize virtual samples at all. For generated virtual contexts e we apply the state regularization method proposed in Section 4.4 for environments with varying state dynamies such Hopper-Mass-OOD and Walker-Mass-OOD.
Here, to obtain on-policy latents Zon for meta-RL, sample Nexp episodes using exploration policy Texp = Te(-|s, z as proposed in Section 4.3. This policy allows to explore a broad range of tasks, as z spans all interpolated areas, including both training tasks and VTs. To enhance the exploration of diverse trajectories, we regenerate z every Hireg timesteps during the exploration process. This periodic update enables the exploration policy to adapt to new tasks and explore the environment with these newly assigned tasks. An amalysis of the exploration policy with respect to the sampling frequency Hireq is provided in Appendix E.2. Finally, summarizing the contents of Section 4 and Appendix D, the proposed TAVT algorithm is outlined in Algorithm 1 (Meta-training of TAVT) and Algorithm 2 (Meta-testing of TAVT).
Algorithm 1 Meta Training of TAVT
Require: The training task set (T.)j=,....NNua. ~ p(Tirain), the task encoder dd. the task decoders Pe and Pæ. the WGAN discriminator fg. the Q-function Qs, and the policy T 1: Initialize parameters th, Φ, c, Ç, and replay buffers for all training tasks.
2: for epoch k 2, do
3: Sample Nmeta training tasks from the training task set.
4: for task Ti in Nmeta training tasks do
5: Collect contexts c using the exploration policy Texp = T[1 z for Nexp episodes.
6: Collect contexts .ff using the RL policy सR.. <-118, 2n for NRI. episodes, where zon qo (1cc)a)
7: Store contexts Ti and off in the replay buffers DJ: and D respectively.
8: end for
9: Construct NvT virtual tasks using the training task latents.
10: Generate the virtual contexts e for each VT using the task decoder Pa.
11:
for gradient step in Model steps do
12:
Sample Netta tasks in train tasks set. Compute Luisim (v), Φ, Φ) loss by Eq. (D.1). Compute Cdise(C) and Egen Φ) losses by Eq. (4) and Eq, (5).
13:
14:
Update the WGAN discriminator parameter by Ç + Ç X1r V (cLissc(().
16:
Update the model parameters yb and Φ, Φ by (v), o, Ø) + (tb, Φb Φ) Aorccontext V ((v.d,)) Frotal 0. δ). end for
17:
for gradient step KRL steps do
19:
Compute RL losses (0) and L (0) by Eq. (D.4) and Eq. (D.5), respectively.
20:
Update the RL parameter A by 0 +- Air Vo(Co(() (+)).
21:
end for
22:
end for
19

---

Task-Aware Virtual Training: Enhancing Generalization in Reinffoeinorreement Learning for Out-of- Tasks
Algorithm 2 Meta Testing of TAVT
Require: The OOD test task set Mtests ne. 9v-
for task T in Muest do
2: for episode k=l, Nexp do
3: Generate zon every Hirg timesteps
4: Collect contexts c.n using the exploration policy Texp - TT| z
5: end for
6: Task inference Zon (-)c)))))
7: Rollout transition using iTa ss Zon) for the last episode.
8: end for
9: Compute the average return of the last episodes for all test tasks.
E. Exploration Trajectories of the Proposed Exploration Policy Texp
In Section 4.3, we propose the exploration policy Texp <(1|s, , z cover a diverse range of tasks, whereas PEARL uses (-1s, 2), ă ((,, I) for exploration. We compare the exploration trajectories of PEARL and our TAVT in Section E. Additionally, in Section D.3, we introduce the method that update z every Hirrq timesteps during the exploration phase to observe diverse trajectories, while PEARL samples ž only once per episode. To demonstrate the effectiveness of Hfreq, we analyze its impact on exploration diversity in Section E.2.
E.1. Trajectory Comparison between PEARL and TAVT
The exploration process in the PEARL algorithm involves sampling from prior distribution NO,, I), as it learns to align all task latents z with this prior, as discussed in Section In contrast, our algorithm uses zo obtained from VT construction, which spans diverse regions the latent space, including both training tasks and virtual tasks. To compare the exploration behavior of PEARL and our TAVT, Fig. E. and Fig. E.2 illustrate first and second exploration trajectories during the exploration phase, as well the final trajectory generated by TRL in the Ant-Goal-OOD environment for training tasks and OOD tasks, respectively, for (a) PEARL and (b) TAVT. From the results, is evident that TAVT's exploration policy covers broader range of areas in the goal space compared to PEARL's exploration policy. Since the task latent for the RL policy is selected by the task encoder based on these exploration trajectories, TAVT's ability to explore diverse goal spaces enhances the diffcrentiation of the current task within the task latent. Consequently, TAVT successfully reaches the goal points for both training and OOD tasks due to the effectiveness of our task latent, whereas PEARL fails to achieve the goal points in both scenarios.
Inner Train
Outer Train
Inner Train
Outer Train
goal
episode 1
y-position
pepssiion
positioon
y- p-position
épisode 2
episade
0
0
0
-2
-2
2
-4
-2
x-position
position
X- x-sostioon
x-
X-position
(a)
(b)
Figure E.I. Visualization of trajectorics the goal space during the meta-testing phase for inner and outer training tasks in the Ant-Goal OOD environment: (a) PEARL and (b) TAVT.
20

---

Task-Aware Virtual Training: Enhancing Generalization in Meinf-eeinooreemeen Learning for Out-of-Distributtin Tasks
OOD
OOD
goal point
Πexp episode 1
2
2
y-position
y-position
Πexp episode 2
episode
-2
-2
-4
-2
-4
-2
X- x-position
x-position
X-
(a)
(b)
Figure E.2. Visualization of trajectories in the goal space during the meta-testing phase for OOD test tasks in the Ant-Goal-OOD
environment: (a) PEARL and (b) TAVT.
E.2. Exploration Trajectories with Respect to Sampling Frequency Htreg
As proposed in Section D.3, we regenerated z every Hireq timesteps for the proposed exploration policy Texp- To demonstrate the effect of Frog on exploration, Fig. E.3 shows the exploration trajectories during the meta-testing phase for different values of Hiqq (a) 5, (b) 20, and (c) 200.
From Fig. E.3, evident that the sampling frequency is too low, such as Hirey 5, the exploration policy results in trajectories that confined to the area around the starting point without covering much distance. In contrast, if the sampling frequency is too high, like Hfreg 200 (the episode horizon for Ant-Goal-OOD environment), the agent reach locations far from the starting point but fails to explore diverse set of directions. We found that balanced sampling frequency, such as Hfreq 20, allows the agent to explore wide range of directions while still covering distant areas effectively.
...... Keyp episode 1
Rexp episode 2
y-position
y-position
y-position
-2
4
-2
2
-4
-4
-2
-4
-2
X- X-position
X- X-position
x-position
(a)
(b)
(c)
Figure E.3. Visualization of exploration trajectories of TAVT according the sampling frequency Havey (a) 5 steps (b) 20 steps (c) 200
steps
21

---

Task-Aware Virtual Training: Enhancing Generalization in Meinffeiino Learning for Out-of-Distrlbation Tasks
F. Hlpeepaarameter Setup for TAVT
this section, we provide a detailed hyperparameter setup for the proposed TAVT. do this, we first define the coefficients for the loss functions related the Bisimulation metric-basse task representation and task-preserving sample generation proposed in Section and Appendix D.I as follows:
Lissim (v), ④) = TTy Abrisim
(IZ 2'al d(Ti, TTPP))
Bisimulation
E. (8,0, Da (F, (n..a,,id (1r (r- 8)2 (5' 5¹)2
Arcon
Reconstruction loss for P3
Arenon B (eeeem Dit & ( (m (r 1)2 (8' 81)2)
Reconstruction loss for
+ off ( zoit) ) on-off latent loss
z' go. (.(.1 T DT. Vi.
(F.6)
(C) AWGAN on (oa) + Eeeepp [onl]]
}]] + Aap Gradient Penalty, (F.7)
Egen(s) (Bo, Φ) = -Eennpol- Prcenpol-Awaan AWGAN fc 2oitt
+rTr Egongo aa)) [ of 19]
(F.8)
WGAN generator loss
preserving
where Abisim is the Bisimulation loss coefficient, Arecon is the reconstruction loss coefficient, λon-off is the on-off loss coefficient, AWGAN is the WGAN loss coefficient, ÀTP the task-preserving loss coefficient, and AçP the cradient-penallee coefficient. Along with these loss coefficients. Table F.1 displays the shared hyperparameter setup across all environments, while Table F.2 and F.3 presents the hyperparameter setup specific to each environment. As shown in Tables F.1, F:2 and F.3 the proposed TAVT has more hyperparameters compared to PEARL. However, most of the loss coefficients are similar across different environments, and the environment-specifn hyperparameters listed in Table F.2 and F.3 are the same as those used in PEARL. Thus, in practice, TAVT does not require a significantly larger hyperparameter search compared to PEARL.
22

---

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcemmee Learning for Out-of-Distributin Tasks
Table F.I. Shared hyperparameters
Name
Value (for MuloCo)
Value (for MetaWorld.
Abisim
100 (50 for Cheetah- (heetah-Vel-OOD)
100
Aeoon
200
200
Aon-off off
Loss coefficients
Awgan
1.0
1.0
preserve
100
100
AVT
1.0 (0.1 for Hopper-Mass-OOn
1.0 (for Reach) / 0.1 (for Push)
Agp
5.0
5.0
Context learning rate contextt
0.0003
0.0003
Learning rate Aιr
0.0003
0.0003
Optimizer
Adam
Adam
Mixing coefficient /3
2.0
2.0
Shared Hyper-
Distance coefficient 17
0.1
1 (for Reach) 10 (for Push)
Regularization coefficient Erog
0.1
parameters
Latent dimension
Batch size for RI.
256
512
Context hatch size No
128
256
Num. of exploration trajectories Nexp
2 (4 for Anl-Goal-OOD)
2
Num. of RL trajectories NRI.
(6 Ant-Dir-4)
Sampling frequency Hnog
20
50
Model gradient steps epoch Kmodel
500 (1000 for Walker/Hopper-Mss--O00
1000
RL gradient steps per epoch KRL
4000 (1000 for Cheetah-Vell Vel-OOD)
4000
4v Qa, me
[300,300,300]
[300,300,300]
Network sizes
Pd. P3 fc
[256,256.256]
[256,256,256]
[200,200,200]
[200,200,200]
Table F.2, MuJoCo environmental hyperparameters
Eavironments
Name
Goal-GGGlOo
Hopper-Mass
Walker
Reward scale Arow Entropy coefficient Aven Num. of training tusks Nirain
5.0
5.0
5.0
1.0
5.0
5.0
0.5
0.5
0.5
0.2
0.2
Environmental Hyperpurameters
100
2
4
150
100
100
Task batch size Naata
16
16
Num. VTs Nvr
2
5
Num. mixining tasks M
2
2
3
3
Table F.3. MetaWorld ML1 environmental hyperparameters
Environments
Name
Reach
Reach-OOD-Inter
Reach-OOD-Exxe
Push
Push-OOD-Inter
Push-OOD-Extra
Reward scule Xvow
1.0
1.0
5.0
Entropy coefficient Aun.
0.2
0.2
0.2
1
Environmental
Num. of training tasks Nusin
50
50
50
50
50
50
Hyperparameters
Task batch size Nuasa
16
16
16
16
Num. of VTs NyT
5
5
5
Num. of mixining tasks M
3
3
3
3
3
23

---

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcemeeee Learning for Out-of-Distributin Tasks
G. More Detailed Experimental Setups
In this section, we provide details about the experiments and experimental setup. G.I, we describe the OOD setup the MuJoCo environments; in G.2, we explain the OOD setup of the MLL environments; in G.3, we outline the baseline setup.
G.I. Mujoco Environments
In this section, we provide a detailed description the computational setup and the MuJoCo environments considered in Section 5. We utilized Mujoco environments from the OpenAl Gym library (Brockman et al., 2016) and employed MuloCo 200 libraries for environments with varying reward functions (Cheetah-Vel-OOD. Ant-Dir-2, Ant-Dir-4, Ant-Goal-OOD), For environments with different state transition dynamics -Hopper-Mass-oee 131 libraries, as suggested in our baseline implementation algorithm, PEARL (Rakelly et al., 2019). In all of Mujoco environments, we used a dense reward setup, which remains the same as the setup used in previous studies (Finn et al., 2017; Rakelly et al., 2019; Zintgraf et al., 2019).
Walker-Mass-OOD)
used MuJoCo
Cheetah- Cheetal-Vel-OOD Vel-OOD
Amt--Dir
Ant-Dir-4
0.5
1.5
Ant- Goal-
Happer SaaperrMass-OOe
Walker-Mass
0.5
3.0
3.0
Figure G. Task configurations for the MuJoCo environments under consideration are presented. Training tasks are indicated in blue, while OOD test tasks are marked red.
To provide more detailed description of the Mujoco environments, Fig. G.1 illustrates the task configurations for the environments considered. In addition, for each Mujoco environment, we explain the design of the reward function, how components of the environment change with different tasks, and the configuration of out-of-distribution OOD tasks as follows:
Cheetah-Vel-OOD- The Cheetah agent tasked with moving at target velocities Vtar where the reward increases as the agent's velocity closely matches the target. The reward function is designed to be the negative value of the difference between the current velocity and the target velocity. For training tasks, target velocities Utar are sampled from the training task space Mtrain [0.0, 0.5) U (3.0, For OOD tasks, target velocities are sampled from the test task space Mtest {0.75, 1.25, 1.75, 2.25, 2.75).
The Ant agent is required to move in target directions fair. with higher rewards given when the agent's movement aligns more accurately with the target direction. The reward function is the dot product of the agent's velocity and the target direction. During training, target directions Oair are sampled from the training task space Murain = (0, 7} For OOD tasks, target directions are sampled from the test task space Miest 1 including extrapolated tasks ( and Za directions) from the training task space.
Ant-Dir-2:
7 }
Ant-Dir-4:
The Ant agent is required to move in target directions Oain similar to the Ant-Dir-2 environment. The key difference lies in the task set. In this environment, the target directions 0air for training tasks are sampled from the training task space Myrain = {0, x
2 म. 35 }. For OOD tasks, the target directions are sampled from the test task space
Mbest 1
24

---

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforeemeen Learning for Out-of-Distribation Tasks
Ant-Goal-OOD: The Ant agent is tasked with reaching specific 2D goal positions given by (Fgoal cos O goal: "goat sin O goat). The agent receives higher reward for reaching the goal position more accu- rately. The goal reward function is designed as the negative L1 distance from the current position to the target goal position. For training, the goal positions Tgool and &goal are sampled from the training task space M with "goal E [0.0, 1.0) U [2.5, 3.0) and Ogoal E [0, 2T.. For OOD test tasks, the goal positions are sampled from the test task space Meest with r "goal = 1.75 and Quaal E {0, π.
2 n. 35 }.
Hopper-Mass-OOLE The Hopper agent is required to move forward while dealing with varying body mass across different tasks, which alters the state transition dynamics. The reward function used is the same as that in the original MuJoCo Hopper environment. The task set is constructed by adjusting the initial mass of all joints of the Hopper using a multiplier macale which is an internal parameter the environment that scales the mass up or down. For training, the body mass multiplier Mscale is sampled from the training task space Murain = (0.0, 0.5) U 330, 3.5). For OOD test tasks, the body mass multiplier is sampled from the test task space Miest (0.75, 1.25, 1.75, 2.25, 2.75).
Walker-Mass-OOn OOD: The Walker2D agent is tasked with moving forward while managing varying body mass across different tasks, which affects the state transition dynamics. The reward function remains consistent with that used in the original MuJoCo Walker2D environment. The task set is created by adjusting the initial mass of all joints of the Walker using a multiplier mecales an internal parameter of the environment that scales the mass either up or down, as similar to Hopper-Mass-On OOD. For training, the body mass multiplier mscale is sampled from the training task space Mtrain [0.0, 0.5) U [3.0, 3.5). For OOD test tasks, the body mass multiplier sampled from the test task space Miest (0.75, 1.25, 1.75, 2,,2,,,,7.
G.2. Meta World ML1 Environments
M
Minner
Shigh
Star (gas 9y: 9z)
Slow
(a)
(b)
(c)
Figure G.2. (a) Reach environment that the end effector should reach to target position. (b) Push environment that the robot arm should push the object the target position. (c) Visualization of MLI-OOD setup that separate 3D goal space M into Minner and M Minner-
ML1-Reach and ML1-Push are environments in the Meta World benchmark (Yu et al., 2020), where a robotic arm either reaches target position in 3D space (MI.1- or pushes an object to a target position (ML.I-Push). Tasks are defined by the goal target position, Itar and their reward functions are determined accordingly. The target position sampled from 3D goal space M, rectangular cuboid defined by two vertices, flow and 9high, which given in the original environment code and differ between -Rel- and M1.h-ush. Training and test tasks are sampled from Merain and Mtest with Ntruin and Ntest tasks, respectively. In the basic MLI setup, MMrrain and Muest are identical to M.
For OOD setups, we introduce a smaller central region, Minner that is defined within M. The goal space M is divided into 125 smaller cuboids (5 x x 5). with Minner consisting of 27 central cuboids (3 x 3 x 3). Fig. G.2(a) represents the Reach environment, Fig. G.2(b) shows the Push environment, and Fig. G.2(c) illustrates the goal configuration in the Meta World environments. Table G.I summarizes the OOD setups for MLI environments.
Basic MI.1 setups: The training goal target positions (train tasks) are sampled with 50 positions from the regions of M, while the testing goal target positions (test tasks) are sampled with 50 positions from the same regions of M.
OOD-Inter setup: The training goal target positions (train tasks) are sampled with 50 positions from the regions of
25

---

Task-Aware Virtual Training: Enhancing Generalization in Mela-Reinforeement Learning for Out-of-Distribation Tasks
M excluding Minner while the testing goal target positions (test tasks) are composed of the center points of the 27 cuboids within Minner-
OOD-Extra setup: The training goal target positions (train tasks) are sampled with 50 positions from within Minners while the testing goal target positions (test tasks) are composed of the center points of the remaining 98 (125 27) cuboids outside Miner.
Table G.J. Configuration of basic MI.1 and ML.1 OOD environments sctup.
Snigh
Mirain
Must
Ntrain
Nuest
Reach
M
M
50
50
Reach-OOD-Inter
(-0.1.08,0..55)
00.1,0.9...)
Minerr
{27 center points Minner }
50
27
Reach-OOD-Eii
Minner
198 center points of MinMneerr]
50
98
Push
M
M
50
50
Push-OOD-Inter
(-0.1,0.8,0.01)
(0.1,0.9..002)
Mininner
27 center points Minner
27
Push-OOD-Extra
Minner
{98 center points of Minnne
50
98
G.3. Baselines Implementation
RI.2 We utilize the open source codebase of LDM at https//gthhh..cmm/uuyuuuggeeeel/ for MuloCo environments and open source codebase garage https://github.com// wowp///wwooggguuuggrgggggg for MetaWorld MLI environments to report the results of RL2. We modify the task space of MuloCo and MI. environments be divided into Mtrain and Muests in other words, OOD setup.
VariBAD and LDM We utilize the open reference of LDM at https//github.comuuuunn---eeee environments and open source reference of SDVT https://ithub..om/uuunnnggg// com/suyoung-lee/SDVT for MetaWorld ML1 environments to acquire the results of VariBAD and LDM. We modify the task space of MuJoCo and MLI environments to be divided into Merain and Meest-
for MuJoCo
MAML We utilize the open source code garage at htgggghthhhoooo////orgrggggggggo for both MuloCo and MI.1 environments to get the results of MAML. We modify the task space MuJoCo and MLL environments be divided into Mirain and Mtest-
PEARL We utilize the open source repository of PEARL
hwt.h//bbcooommaaaaarelllllyyeee for MuJoCo environments for MLI environments to get the results of
and the open source repository of garage PEARL. We modify the task space of MuJoCo and ML1 environments to be divided into Mirain and Muest-
CCM, MIER, Amago We utilize the open source code of CCM, MIER and Amago at https://ittuu..com/TIi LAB/self wwwssiiieedr///eeeeeeeeeeeeeeeee 01t.::2221..61404499994477791559 https://github..c http://...oo.eeeleenennnn//////e httpoonnneaeeeeaieeeiieeeea and https://github..cmm Sttp://T--ssti-RP--LLPP
b8c4981
with Environment RetpememenReeeeeetttttoncccemm
respectively for both MuJoCo and MLI environments to measure the experimental results of the baselines. We modify the task space of MuJoCo and ML1 environments to be divided into Mtruin and M
TAVT We research and develop TAVT algorithm on top of the PEARL official open-source algorithm at for both MuJoCo environments and MLI environments. Our implementation code is
available at
We the MetaWorld benchmark open source at http::/github.comm install the "paper version" of it. All experiments conducted GPU server with an NVIDIA GeForce RTX 3090 GPU and AMD EPYC 7513 32-Core processors running Ubuntu 20.04, using PyTorc..
for all experiments, but
26

---

Task-Aware Virtual Training: Enhancing Generalization in Meta- Learning for Out-of-Distribation Tasks
H. Additional Results for TAVT
H.1. Learning Curves for Meta World ML1 Environments
Fig. H.l presents the learning curves of the MLL environments from the comparison experiments in Section 5. Similar to the results in Table 3, the learning curves show that the proposed TAVT consistently outperforms other on-policy and off-policy meta-RL methods. Notably, compared off-policy algorithms, TAVT demonstrates both superior convergence performance and faster convergence speed. These results further highlight the effectiveness and superiority of the proposed TAVT algorithm.
1.0
1.0
1.0
Success Rate
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.
Time Steps 1e7
Time Steps 1e7
Time Steps 1e7
(a) Reach
(b) Reach-OOD-Inter
(c) Reach-OOD-Extra
1.0
1.0
1.0
Success Rate
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0.0
0.2
0.4
0.6
0.8 10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time Steps 1e7
Time Steps 1e7
Time Steps 1e7
(d) Push
(e) Push-OOD-Inter
(f) Push-OOD-Extra
TAVT
VariBAD
CCM
PEARL
Amago
MAML
RL?
LDM
MIER
Figure H.I. (a-f) Performance comparison of the 6 MetaWorld MI.1 environments
H.2. Computational Cost
We provide a breakdown the relative computational cost introduced by each component in TAVT in the same setup. We report in Table H.1 the relative training time per epoch that occurs when each component of TAVT is added compared to the PEARL algorithm we base our implementation on. shown in Table H. adding the reconstruction loss and decoder contributes approximately 3% overhead compared PEARL, while the metric-based representation (w/o Lgen) adds about 11%. While full TAVT requires approximately 18% more training time than the PEARL baseline due to the additional training of each component, other algorithms do not yield comparable performance gains even with similar computational overhead, demonstrating the practical advantage of TAVT. TAVT significantly outperforms all baselines across diverse OOD environments in MuJoCo and Meta Morld as shown in Fig. 7. Even with extended training, other methods do not reach TAVT's level generalization. We therefore believe the added cost is reasonable and practical.
Table H.J. Relative per-epach training time comparison over component variants
Method
PEARI.
Recon only
TAVT w/o Eyon
TAVT w/o on-off loss
TAVT(full)
Relative Training Time per Epoch
100%
103%
111%
116%
118%
27

---

Task-Aware Virtual Training: Enhancing Generalization in Meia-Reinforeementt Learning for Out-of-Distribation Tasks
I. Ablation Study on the Mixing Coefficient /
In this section, we discuss the impact the mixing coefficient 3 for VT construction introduced in Section 2. Recall that the task latent for VT, denoted z", is obtained by interpolating among M randomly sampled training task latents, follows:
M
z
a'z'
where (a'!.... . 9iriichee(11,111 1,.... 1) the Dirichlet distribution, and 3 2 1 is the mixing coefficient. When 3 = 1, the interpolation is limited to within the space of the training task latents. In contrast, 3 >> 1 enables extrapolation beyond the original task latents, allowing for broader range of virtual tasks to be generated. To analyze the impact of the mixing coefficient 3 on VT construction and learning performance, we examine the coverage range of 2 with varying / Section I.1, demostrate how exploration behavior changes with different 33 values in Section 1.2, and evaluate the performance based on different /3 values in Section I.3.
8-1 is the interpolation coefficient. Here, Dirichlet(:) represents
M
L.1. Coverage of z with Various
(1
(8=3)
Triniig tasks 000 tasks
Training tasks 0OD tasks
Training tasks 0OD tasks
(a)
(b)
(c)
(d)
Figure 1.1. (a) 2D goal positions of the Ant-Goal-OOD environment, and PCA visualizations of z" with diffcrent values: (b) TAVT
with 8 = 1, (c) TAVT with 2, and (d) TAVT with 3 3.
Fig. I.1 illustrates how the coverage of task latents z for VTs generated from the Dirichlet distribution varies with the mixing coefficient
the Ant-Goal-OOD environment. Specifically, Fig. (a) displays the 2D goal positions for training and OOD tasks, while Figs. I.1(b), (c), and (d) show the principal component analysis (PCA) visualizations of z" generated with / 1, 2, and 3, respectively. Here, data structure in the projection, whereas t-SNE often clusters data by performing manifold learning. The results in Fig. L.1 show that as 3 increases, the range of task latents n generated by VT construction becomes broader, effectively covering a wider area, including extrapolated task latents as intended in Section 2.
PCA instead of t-SNE for visualization, PCA more accurately reflects the
1.2. Exploration Trajectories of Texp with Various 3
Fig. 1.2 illustrates how exploration trajectories change with different 3 values during the meta-testing phase. For 3 = 1 in Fig. I.2(a), z considers only a narrow range of the task latent space, causing the agent to explore only the areas close to the inner region. For / = in Fig. I.2(b), zd covers a broader range, which is sufficient to distinguish the tasks from the contexts. For 3 3 in Fig. I.2(c), z covers the largest area but includes excessive extrapolated regions, causing the agent explore beyond the goal space and making task distinction more challenging. Thus, we use 3 2 as the default mixing coefficient for TAVT.
28

---

Task-Aware Virtual Training: Enhancing Generalization in Reda-Reeiforeemeent Learning for Out-of-Distrlaatio Tasks
Rex episode
y-position
y-position
2
y-position
Nex episode
-4
-2
x-position
X-position
X-position
(a)
(b)
(c)
Figure 1.2. Exploration trajectories of Fexp
(-18, (exp(118, 2²) with different (3: (a) 1, and (c) 3.
L.3. Performance Comparison with Various 3
From the results in Sections I.1 and L.2, we observed that the choice 3 significantly impacts TAVT learning. Finally, to evaluate the effect on performance, Fig. I.3 illustrates the performance changes across 3 environments, (a) Cheetah-Ve Vel- OOD, (b) Goal-OOD and (c) Walker-Man Mass- OOD, with varying values. From the results in Fig. I.3, we observe that consistently performs the best across the considered environments. Therefore, set B the default mixing coefficient for TAVT.
-100
Test Avg Return
Test Avg Return
-150
Test Avg Return
-150
-200
500
-200
-250
400
-250
-300
B 1
300
B=1
-300
==
-350
B= 2
200
8 2
==3 3
==3
-350
-400
100
0.5
1.0
Time Steps
le6
Time Steps
1e6
Time Steps
1e6
(a)
(b)
(C)
Figure 1.3. Performance comparison with various values on: (a) Cheetah-Vel Vel-OOD environment, (b) Ant-Goal-OOD environment, and (c) Walker-Mass-OOn environment.
29


Reviews:
[{"review_id": "83Q533X3yM", "reviewer": "Reviewer_vcob", "summary": "This paper provides a new method for meta Reinforcement Learning that focuses on generalizing to novel tasks. The paper provides a method for generating virtual tasks that takes in to account task characteristics. A metric-based task representation is learned using the Bisimulation metric. Furthermore task preserving sample generation is used to generate additional samples generalize to unseen tasks. With a combination of these, Task-Aware Virtual Training, trains a task latent dependent policy using Soft-Actor Critic. The paper compares the new method against several baselines across many different environments.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thoughtful and constructive comments. We would like to address the main concerns raised, particularly regarding (1) issues of clarity, and (2) technical questions.\n\n**1. Clarity Improvements in Weaknesses Part:**\n\nWe sincerely thank the reviewer for pointing out aspects of the paper that could benefit from improved clarity.\n\n- Section 4.3 Algorithm Description: We appreciate the feedback that the algorithmic workflow in Section 4.3 may be difficult to follow in its current narrative form. To improve readability, we will revise this section to include a clearer step-by-step breakdown, either by adding a numbered list or directly referencing steps from Algorithm 1 in the appendix.\n\n- Limitations in Appendix: We agree that the discussion of limitations is important and should be accessible to all readers. We will relocate the key points from the appendix into the main body of the paper, allowing them to be more visible and impactful.\n\n**2. Technical Clarifications for Questions:**\n\nThank you for the thoughtful questions. Below are our clarifications, which will also be reflected in the revised manuscript:\n\n- Averaging in Figure 6 and Table 3: All performance metrics in these figures are averaged over 5 random seeds. For each seed, the reported value reflects the average performance across all considered OOD test tasks in that environment. The shaded areas and error bars represent the standard deviation across these seeds. We will explicitly clarify this in both the main text and figure captions to avoid any ambiguity.\n\n- Equations D.4 and D.5 (Source of $\\mathbf{z}^\\alpha$ and $\\hat{\\mathbf{c}}^\\alpha$): The virtual latent $\\mathbf{z}^\\alpha$ is computed through linear interpolation between training task latents as described in Section 2.3. The virtual context $\\hat{\\mathbf{c}}^\\alpha$ is then generated using the task decoder conditioned on $\\mathbf{z}^\\alpha$ as described in Section 4.2. We will revise the appendix to clearly connect these definitions across sections.\n\n- \"First, generated VTs often fail to capture task characteristics accurately\": We would like to clarify that the statement was not referencing a claim made in prior literature, but rather reflects our own empirical observation based on the results shown representation of LDM in Figure 1. Specifically, although LDM also utilizes VT generation, we observe that its task representations do not align well with task semantics such as goal direction in the Ant-Goal environment. We will revise the manuscript to clearly indicate this.\n\n- Figure 1 (1)→(8) Annotation: The indices (1) through (8) represent different goal directions in the Ant-Goal environment as defined in Figure 1(a). These were included in Figure 1(e) to illustrate that our method accurately reflects structured goal alignment. To improve comparison, we will add the same labeling to Figures 1(b–d) and enhance the caption to clarify this design.\n\n- \"allows extrapolation beyond the original latents\": This refers to interpolating task latents using coefficients $\\alpha$ beyond the $[0,1]$ range, effectively generating virtual tasks that lie outside the convex hull of the training task distribution. This helps the agent learn from extrapolated (OOD) regions.\n\n- \"[...] virtual contexts generalize to unseen tasks\": This sentence explains the rationale for generating virtual contexts. Since real samples from unseen tasks are unavailable, the generated contexts do not correspond to any specific real task. As the reviewer correctly inferred, this is closer to the latter interpretation, virtual contexts are generated by decoding linearly interpolated task latents using the decoder’s generalization capability. We will clarify this point in the revised version.\n\n**3. Other Comments and Typos:**\n\nWe are grateful for the reviewer’s careful reading. The following minor issues will be addressed in the revised manuscript:\n\n- Expand \"LDM\" to \"Latent Dynamics Mixture\" at first mention.\n- Correct phrasing such as “identify several key issues” to \"identify two key issues.\"\n- Fix typos like \"aallows\" and remove unnecessary commas.\n- Introduce the term \"task decoder\" earlier in Section 2.2 for clarity.\n\nOnce again, we sincerely thank the reviewer for their constructive feedback and positive recommendation. We believe these clarifications will further improve the clarity and completeness of our work.", "claims_and_evidence": "- A claim is made that the task their task-preserving loss ensures that the virtual contexts better retain task latent information. If the meaning of this is that there is parity between the virtual latent and the output of the encoder to the virtual context, then I believe that Figure 3 makes a reasonable case for this claim. \n\n- A claim is made that their exploration policy covers a _wider_ range of trajectories. This was compared with PEARL. The evidence in Appendix E shows for the Ant-Goal problem what trajectories from exploration policies trajectories look like after the first and second episodes looked like. I think that this claim is visually shown in the Figures. \n\n- A claim is made that the proposed method out-performs the compared baselines, as well as that generalization is good. I believe that the results from Figure 6 and Table 3 show case this.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I checked the Proof of Proposition 3.1. I believe it is correct.", "experimental_designs_or_analyses": "- It is unclear to me how many repeats of each experiment were carried out to obtain the results in Figure 6 and Table 3. \n\n- Am I right in understanding that the averages in Figure 6 for example are also across the different test tasks (for each environment)?", "relation_to_broader_scientific_literature": "This paper does a good literature review, and also carries out comparisons to other meta learning algorithms. The paper builds on top of methods from Latent Dynamic Mixtures [1], and also uses the Bisimulation metric [2]\n\n[1] Lee et al, Improving generalization in meta-rl with imaginary tasks from latent dynamics mixture. 2021\n[2] Zhang et al, Learning invariant representations for reinforcement learning without reconstruction. 2021", "essential_references_not_discussed": "None to my knowledge.", "other_strengths_and_weaknesses": "### Strengths\n- Good literature review\n- Generally written well\n- Good experimental comparisons across many different method from the literature.\n- Appendix F and G details a lot of the experimental setup. I believe this is sufficient for reproducibility.\n- Good ablation studies are carried out. \n### Weaknesses\n- I had trouble understanding the algorithm described in words in Section 4.3 until I looked at the Algorithm 1 in the Appendix. Perhaps the steps of the algorithm can be explained step-by-step (by breaking into a numbered list for example). \n- Limitations are discussed in the Appendix. Perhaps this should be done in the main text.", "other_comments_or_suggestions": "### Introduction\n1. LDM is used before writing its expansion\n2. Figure 1: Can a description of what the annotation (1) -> (8) represents in Figure 1(e) be written in the caption.\n3. 'Despite its benefits, we identify _several_ key issues with the existing methods for VT construction' -> Since you identify two issues, wouldn't it be more prudent to say '[...] identify _two_ key issues [...].'?\n\n### Section 2\n4. After the definition of the task encoder loss, there is an extra comma.\n5. End of Page 2: typo 'aallows'\n6. End of Page 2: You mention a 'task decoder' for the first time as if it was mentioned previously in the section. Could you describe the 'task decoder' in the first paragraph of Section 2.2 please?", "questions_for_authors": "### Introduction\n1. 'First, generated VTs often fail to capture task characteristics accurately.' -> What is the justification for this? Is this from your own analysis or from the literature. If the latter, please give a reference.\n2. Figure 1:  Is there a notion in which there would be equivalent points (1) -> (8) in (b) - (d)?\n\n### Section 2\n3. What do you mean by 'allows extrapolation beyond the original latents'? \n\n#### Section 4\n4. line 193: What do you mean by '[...] virtual contexts generalize to unseen tasks.'. Do you mean that the generated contexts would correspond to real 'unseen' tasks? Or that the task latent that would correspond to the generated contexts would be close to virtual context?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RY4KMcNbfE", "reviewer": "Reviewer_4UJT", "summary": "This paper introduces Task-Aware Virtual Training (TAVT), a novel meta-reinforcement learning algorithm designed to enhance generalization to out-of-distribution (OOD) tasks.  The numerical results demonstrate that TAVT significantly enhances generalization to OOD tasks across various MuJoCo and Meta-World environments.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thoughtful and constructive comments. We would like to address the main concerns raised, particularly regarding (1) the novelty of our method, (2) benchmark coverage, and (3) computational cost.\n\n**1. Novelty of Our Approach:**\n\nWe appreciate the reviewer’s comment on novelty. To address the concern, we clarify that our method is not a generic combination of task representation and VT, but a principled framework designed to answer a central question: **how can we generate virtual tasks that are truly task-aware**, that is, faithful to underlying task dynamics and semantics?\n\nTo this end, we propose a metric-based task representation that aligns task latents using the Bisimulation metric. Unlike contrastive methods that separate positive and negative pairs without modeling precise distances, the Bisimulation metric captures behavioral similarity through reward and transition dynamics. This enables the latent space to reflect task geometry, supporting meaningful interpolation and extrapolation. We also introduce an on-off alignment loss to enforce consistency between off-policy latents (used in training) and on-policy latents (used in meta-testing), which is essential for stable generalization. As shown in Fig. 1, our method yields more structured and consistent task representations compared to contrastive-learning-based methods like CCM.\n\nBuilding on this aligned space, we present a task-preserving sample generation strategy using WGAN, coupled with a latent consistency constraint: generated virtual contexts should reconstruct their source latent when re-encoded. This ensures that generated data remains both realistic and task-consistent. As demonstrated in Fig. 3, our method effectively captures OOD task semantics, in contrast to prior VT-based methods that suffer from latent mismatch.\n\nTogether, these two components form the core of TAVT. We believe this task-aware design is the key to its strong OOD generalization, supported by our visualizations and ablations. While representation learning is a widely explored area in meta-RL, it remains a crucial pillar for generalization. Our work extends this line by explicitly aligning representations with task semantics during both training and deployment.\n\n**2. Benchmark Scope:**\nWe appreciate the reviewer’s thoughtful suggestion regarding ML10 and ML45. These benchmarks feature nonparametric task variation (e.g., changes in object types or task semantics), which differ fundamentally from the parametric setups we target, such as variations in velocity, mass, or 3D goal positions of tasks. In parametric cases, virtual tasks can be generated meaningfully through latent interpolation, making VT-based augmentation both tractable and effective. In contrast, ML10/45 often require task decomposition or modular policies due to their discrete and heterogeneous structure. \n\nFor example, SDVT [1] tackles this by clustering tasks through task decomposition and applying per-cluster VT modules using LDM, rather than building a unified latent space. Accordingly, our method can be seen as complementary: integrating our task-aware VT generation within each sub-task cluster in nonparametric settings could enhance generalization. We will include this as a promising future direction in the revised version.\n\n[1] Lee et al., \"Parameterizing non-parametric meta-reinforcement learning tasks via subtask decomposition,\" NeurIPS 2023.\n\n**3. Computational Cost:** (common response with Reviewer t99j)\n\nWe appreciate the reviewers’ concern regarding the potential computational overhead of TAVT. While it is difficult to directly compare training time across all on-policy and off-policy baselines, we provide a breakdown of the relative cost introduced by each component in TAVT in the same setup. As shown in Table R.1, adding the reconstruction loss and decoder contributes approximately 3% overhead compared to PEARL, while the metric-based representation (w/o $\\mathcal{L}_{\\text{gen}}$) adds about 11%. When the full TAVT is considered (using WGAN), the total overhead reaches approximately 18%.\n\nDespite this increase, TAVT significantly outperforms all baselines across diverse OOD environments in MuJoCo and MetaWorld, as shown in Fig. 6. Even with extended training, other methods do not reach TAVT’s level of generalization. We therefore believe the added cost is reasonable and practical. This analysis will be added to the revised version for improved clarity.\n\n**Table R.1:** Relative Per-Epoch Training Time Across Component Variants\n|Method|Relative Training Time per Epoch|\n|-|-|\n|PEARL|100%|\n|Recon only|103%|\n|TAVT w/o $\\mathcal{L}_{\\text{gen}}$|111%|\n|TAVT w/o on-off loss|116%|\n|TAVT (full)|118%|\n\nOnce again, we thank the reviewer for their detailed feedback and insightful suggestions. We hope our clarifications help address the raised concerns and further highlight the contribution of our work within the broader meta-RL literature.", "claims_and_evidence": "The paper claims that TAVT enhances generalization to OOD tasks.  This claim is supported by numerical results in MuJoCo and Meta-World environments.  The results consistently show TAVT outperforming other meta-RL methods in OOD settings.", "methods_and_evaluation_criteria": "The method applies several mechanisms, including metric-based task representation, data generation, and state regularization to improve the OOD generalization ability of context-based RL. The proposed methods somehow make sense.\n\nThe evaluation criteria include comparing TAVT with various on-policy and off-policy meta-RL methods across MuJoCo and MetaWorld environments.  The performance is measured by the average return on OOD test tasks.", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "The experimental design involves comparing TAVT with other meta-RL methods in MuJoCo and MetaWorld environments.  The authors analyze the task representation of TAVT and conduct ablation studies to evaluate the contribution of different components of TAVT.", "relation_to_broader_scientific_literature": "The method applies several mechanisms, including metric-based task representation, data generation, and state regularization to improve the OOD generalization ability of context-based RL.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths\n\nThe paper is well-structured and clearly written.\n\nThe paper addresses an important problem in meta-RL: generalization to OOD tasks.\n\nThe paper proposes a novel and effective algorithm, TAVT, that outperforms existing methods.\n\nThe paper provides comprehensive experimental results and analysis to support its claims.  \n\nWeaknesses\n\nThe method applies several mechanisms, including metric-based task representation, data generation, and state regularization to improve the OOD generalization ability of context-based RL. The proposed methods somehow make sense. However, the mechanisms are not novel.  For example, the data generation and data mix-up are widely applied to improve the generalization. The bisim loss in metric-based task representation is similar to the combination of the contrastive loss and reconstruction loss, which has been widely discussed in context-based RL papers.\n\nThe proposed method generates virtual task samples using GAN, which may increase the training cost of context-based RL.", "other_comments_or_suggestions": "MetaWorld is a good benchmark with OOD test tasks. The most important feature is the heterogeneous types of tasks in ML10 and ML45. However, these scenarios are not tested by the paper.", "questions_for_authors": "In the bisim loss (3), is the computation cost of the task distance metric high?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BJQWJd7oDo", "reviewer": "Reviewer_q6Z7", "summary": "In context-based meta RL, an agent is trained to assign a context vector to an MDP. It learns to identify the context, and to take actions (policy) based on it. This paper investigates the setting where an agent is tested in an MDP not seen during training.\nSeveral ideas are introduced to improve generalization to unseen MDPs:\n1. Using a bisimulation metric idea, contexts of bi-similar MDPs are trained to be similar.\n2. The agent is trained on imagined data from contexts that are \"mixes\" of seen contexts. \n3. A GAN is used to learn better reward+next state predictions for the imagined contexts.\nIn experiments, agents are evaluated on MDPs that were not trained on. The variations in the MDPs are all  low-dimensional - a single velocity direction, a 2/3 dimensional goal, a single mass coefficient, etc. The proposed method is shown to outperform previous methods.\n\n## Post rebuttal:\nThe authors' rebuttal clarified the comaprison with Mendonca et al., 2020 (MIER) and mostly addressed my main concern of comparison with baselines tailored for OOD meta RL. From the list of references the authors mention, Ajay et al., 2022 also seems like a relevant (and more recent) comparison, but the paper already includes a rather extensive experimental evaluation.\n\nOverall, and after reading the other reviews, I'm more convinced that this is an incremental but solid contribution to OOD tasks in Meta RL, and is a good fit for ICML.\n\nI have changed my score to weak accept.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thoughtful and constructive comments. We would like to address the main concerns raised, particularly regarding (1) comparison with baselines handling distributional shift, (2) the scope and practicality of our method, (3) clarification on reward functions, (4) the connection to theoretical work on generalization, and (5) additional future work.\n\n**1. Comparison with Baselines Addressing Distributional Shift:**\n\nWe appreciate the reviewer’s insightful suggestion to compare with prior work that explicitly targets distributional shift. We would like to clarify that among the baselines the reviewer mentioned, MIER (Mendonca et al., 2020) is indeed already included in our empirical evaluation. We selected MIER as it is a PEARL-based method specifically designed to handle distributional shift, closely aligned with our setting, and it provides open-source code for reproducibility. Additionally, as the reviewer noted, LDM is also an algorithm that addresses OOD test tasks through virtual task generation, and is therefore included in our comparative evaluation.  We regret that the citation for Mendonca et al. (2020) was inadvertently omitted, and we will correct this in the revised version to avoid confusion.\n\nMoreover, we designed our comparison to cover a diverse range of meta-RL algorithms across categories, including both on-policy and off-policy methods (e.g., VariBad, RL$^2$, PEARL), methods with advanced task representations (e.g., CCM using contrastive learning, Amago using transformer-based in-context learning), as well as those employing virtual tasks (e.g., LDM). As a result, we focused on representative or most relevant algorithms per category rather than exhaustively including all prior works, due to experimental feasibility.\n\n**2. Scope and Practicality of the Method:**\n\nWe acknowledge and appreciate the reviewer’s comment on the scope of our work. While our approach lies within the context of meta-RL with virtual task training, we argue that it has broader practical implications. One of the core contributions of TAVT is its ability to generalize across a wide spectrum of task shifts, including: \nState transition variations (e.g., in Walker-Mass/Hopper-Mass), reward variations (e.g., Ant-Dir/Cheetah-Vel), and 3D task variation (ML1-Push/Reach).\n\nBy addressing these diverse sources of task variability, TAVT takes a step closer toward real-world applicability, where test tasks often differ substantially from those seen during training. We agree that practical validation in real robotic systems or high-dimensional continuous control domains remains a promising future direction, and we will add a discussion of this in the revised version.\n\n**3. Clarification on Reward Functions:**\n\nThank you for highlighting the need to clarify the reward functions used in our experiments. We confirm that dense rewards were used in all environments, following standard practice in recent meta-RL literature. For example, Ant-Goal-OOD uses dense distance-based rewards as in PEARL and LDM. We will explicitly clarify these reward formulations in the main text and provide detailed specifications in Appendix G to ensure transparency.\n\n**4. Connection to Theoretical Work on Generalization:**\n\nWe thank the reviewer for pointing out the connection to Rimon et al. (2022), which presents a theoretical framework for understanding generalization in meta-RL via density estimation. Their work highlights the challenge of generalizing to low-density regions in the task distribution. In contrast, our method complements this by densifying the latent task space through virtual task generation and bisimulation-aligned interpolation. This allows the agent to train on transitions that approximate those in unseen test tasks.\n\nWe see this connection as a valuable opportunity for synergy: combining our sample generation strategy with density-aware task weighting could lead to more principled and efficient task coverage, especially under limited task availability. We will cite and discuss this work in the revised paper and consider integrating its ideas in future research.\n\n**5. Additional Future Work:**\n\nWe appreciate the reviewer’s encouragement to explore broader extensions. Beyond the theoretical connection mentioned above, one promising direction is the application of TAVT to offline meta-RL. In offline settings, OOD task generalization is especially challenging due to limited and biased data. We believe that our VT-based sample generation can augment offline datasets with behaviorally consistent transitions, helping mitigate the distributional gap between training and test tasks. We plan to explore this extension in future work.\n\nOnce again, we thank the reviewer for their detailed feedback and thoughtful suggestions. We hope our clarifications address the raised concerns and help situate our contribution within the broader meta-RL literature.", "claims_and_evidence": "The experiments are thorough, and the paper is well written. The main idea makes sense, and extensive experiments support it (but see my question about missing baselines later).\n\nOne issue that should be emphasized: the task distributions here are all low-dimensional, and I suspect that this is a limitation of the method (like it is for other context-based meta RL).", "methods_and_evaluation_criteria": "The authors experiment with several benchmark tasks that are standard in meta RL, and the comaprison makes sense. \n\nIt would be good if the authors explicitly discuss the reward functions used. I had to look in the code to verify that dense rewards were used. Dense rewards make some domain (like ant goal) much easier, but this is a choice that has been popularized in several prior works.\n\nI was missing baselines that explicitly account for distribution shift, except for LDM. Citing from the related work section: \"some studies tackle distributional shift challenges through robust learning (Mendonca et al., 2020; Mehta et al., 2020; Ajay et al., 2022; Greenberg et al., 2023)\". Why are these methods not compared with?", "theoretical_claims": "The (fairly simple) claim on the bi-simulation metric being a metric looks correct.", "experimental_designs_or_analyses": "In high level the experiments follow standard meta RL evaluations and look OK.", "relation_to_broader_scientific_literature": "The paper contribution is relatively narrow: within the scope of meta-RL with low-dimensional task distribution, it proposes several technical advances for improving generalization. The idea of training on virtual tasks has been proposed before, but the technical improvements proposed here make sense and are reported to work well. \n\nThe authors nicely discuss other meta RL studies on OOD generalization in the related owrk section, but these are mostly missing from the empirical evaluation.", "essential_references_not_discussed": "A relevant theoretical investigation of generalization in meta RL, and how virtual tasks can help generalization appears in [1] and should be mentioned. That paper also explicitly discussed the dimension of the task space, and why such methods are limited to low dimensional tasks.\n\n[1] Rimon, Z., Tamar, A., & Adler, G. (2022). Meta reinforcement learning with finite training tasks-a density estimation approach. Advances in Neural Information Processing Systems, 35, 13640-13653.", "other_strengths_and_weaknesses": "The paper is very well written, easy to follow, and the appendix provides extensive details. Also, code is provided (although I did not run it).\n\nThis paper has a rather narrow scope: how to technically improve meta RL by improving the virtual tasks the agent is trained on. Within this scope, the authors have done a good job at coming up with several good ideas, and extensively evaluated their approach. The overall method is relatively complex (includes training a GAN, several loss terms, etc), but since the authors provide code, I guess that only time will tell if it gets picked up and used in more practical applications (beyond the toy benchmarks).\n\nOverall, if it wasn't for my reservations about the missing baselines, I would have proposed to accept this paper, as it makes an incremental but clear contribution to meta RL.", "other_comments_or_suggestions": "It may be worthwhile to discuss some future work. How can some of the ideas here be used beyond this specific application, or what problems are still missing in the current solution.", "questions_for_authors": "Why is there no empirical comparison with methods that explicitly account for distribution shift?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3VRBQ6rbMZ", "reviewer": "Reviewer_t99j", "summary": "This paper tackles meta RL, where the goal is find optimal policies for MDPs that have the same state-action space but different transition and reward functions. This paper is similar to prior works like PEARL but additionally uses extra loss terms/objectives to improve performance. Their main ideas are:\n- Virtual tasks generation as linear combinations of observed training tasks\n- They propose a distance metric for tasks and train their latent representations of tasks to reflect this distance (bisimulation loss)\n- They train latents for the same tasks to be consistent whether the policy is exploratory on optimizing for that task (on-off)\n- Reconstruction loss for generated samples\n- Task Preserving sample generation: They train the encoder-decoder to be consistent in the sense that the generated samples from some latent z^alpha are encoded to some value hat(z)^alpha that is close to z^alpha. \n- They use a GAN based approach to ensure the generated samples are indistinguishable from real samples\n- They regularize state generation with some real next states\n- They train a policy on both real and virtual tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thoughtful and constructive comments. We would like to address the main concerns raised, particularly regarding (1) computational cost, (2) the validity of state regularizaion method, and (3) t-SNE -based visualization.\n\n**1. Computational Cost:** (common response with Reviewer q6Z7)\n\nWe appreciate the reviewer’s concern regarding the potential overhead of TAVT. While direct comparison of training time across all baselines is difficult  (e.g., training time on the x-axis), we analyze the relative overhead of each TAVT component in a consistent setup. As this is a shared concern and in line with ICML’s policy encouraging shared responses, we kindly refer to **our detailed answer to Reviewer q6Z7**, which includes a full analysis. We thank for your understanding and believe it adequately addresses the concern.\n\n**2. Validity and Sensitivity of the State Regularization Method:**\n\nWe agree that linear state blending may be problematic in discrete or non-Euclidean environments (e.g., image-based or maze tasks). In such cases, an alternative is to use a discriminator-guided filtering strategy, as in recent data augmentation works [1], to retain only realistic generated states. This would improve applicability to more complex domains and is a promising direction for future work. In our study, however, we focused on continuous settings, where such blending is a reasonable approximation due to smooth and Euclidean state dynamics. We will also make sure to clarify this assumption in the revised version.\n\nIn addition, to address the reviewer’s insightful comment on $\\epsilon_{\\mathrm{reg}}$ sensitivity, we extended our ablation study to include a finer sweep over $\\epsilon_{\\mathrm{reg}} \\in [0.0, 0.05, 0.1, 0.2, 0.5, 1.0]$ on the Walker-Mass-OOD environment. The results, shown in Table R.2, confirm that small values around $0.1$ consistently lead to lower $Q$-function estimation bias and improved return, while both extremes ($0.0$ or $1.0$) result in performance degradation. This suggests a convex relationship and supports the reviewer’s observation that a regularization sweet spot exists.\n\n**Table R.2:** Ablation on $\\epsilon_{\\mathrm{reg}}$ in Walker-Mass-OOD\n|$\\epsilon_{\\mathrm{reg}}$|0.0|0.05|0.1|0.2|0.5|1.0|\n|-|-|-|-|-|-|-|\n|Avg. Q Bias|-24.1±25.6|91.4±21.2|**36.4**±23.5|65.1±26.8|98.4±36.4|73.1±32.6|\n|Avg. Return|612.6±27.2|660.8±23.4|**745.6**±42.3|632.3±91.7|238.2±98.3|163.8±96.4|\n\n[1] Zhang et al., \"Discriminator-guided model-based offline imitation learning,\" CoRL 2023.\n\n**3. Use of t-SNE and Latent Space Visualization:**\n\nWe thank the reviewer for the valuable suggestion regarding the limitations of t-SNE and the use of alternative visualization techniques. While t-SNE is a common tool for qualitatively assessing representation structure, we agree that it can be sensitive to hyperparameters and may not always preserve the true geometry of the latent space.\n\nTo address this, we additionally visualized task latents by directly setting the latent dimension to 2 in the Ant-Goal-OOD environment, thereby avoiding t-SNE entirely. The resulting plots (linked below) show raw 2D latent representations without projection, and the alignment with task goal positions confirms that our method preserves semantic structure even without dimensionality reduction artifacts.\n\nLink: https://anonymous.4open.science/r/024/ant-goal-ood-2D.PNG\n\nWe will include these results in the revised version and appreciate the reviewer’s suggestion, which strengthened the empirical support for our metric-based task representation method.\n\n**4. Minor Points and Additional Questions:**\n\nWe sincerely thank the reviewer for their detailed reading and thoughtful suggestions. We address the raised points as follows and will incorporate the necessary revisions in the updated version:\n\n- Proposition 3.1: While the theoretical metric is defined via true expectations, our practical implementation relies on samples from learned transition models. As such, minor approximation errors may occur, particularly in unvisited regions of the state space across tasks. We will clarify this limitation.\n\n- Typos and notation issues: We appreciate the reviewer’s careful attention and will correct all identified typos and symbols in the revision.\n\n- Figure 6: On-policy methods generally require significantly more samples for convergence. Following common practice in meta-RL, we report their final performance when comparing with off-policy methods.\n\n- One-hot encoding: The one-hot decoder is used only for training tasks. While scalability could be a concern with an extremely large number of training tasks (usually not large due to memory constraints), this has no impact on test-time generalization.\n\nOnce again, we sincerely thank the reviewer for their constructive feedback and positive recommendation. We believe these clarifications will further improve the clarity and completeness of our work.", "claims_and_evidence": "The claims are effectively that all of the above techniques are necessary and performant. Their claims are empirically validated, where ablations show that performance significantly degrades if any technique is removed.", "methods_and_evaluation_criteria": "Yes, the proposed approach and techniques are well-motivated given the problem.", "theoretical_claims": "I've checked the theoretical claims/proof in appendix A. It appears correct, but I believe there are some almost everywhere arguments that need to be addressed for completeness. The proposed metric may be a pseudo metric, though this has little practical implication.", "experimental_designs_or_analyses": "The experiments are standard for Meta RL settings. Using prior implementations for baselines is good practice.", "relation_to_broader_scientific_literature": "This paper introduces novel contributions to prior meta RL algorithms, and I consider it an improved version of PEARL.  It also leverages ideas from GANS, VAEs, bisimulation, etc.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n- The method is thoroughly ablated\n- The empirical results are on reasonably challenging benchmarks\n- Each individual contribution is well-justified.  \n\nWeaknesses:\n- This paper does not sufficiently address its limitations. Primarily, I am concerned about compute time. This paper introduces many new networks and loss terms, which greatly increases resource usage. In the appendix, the authors claim that their algorithms has better convergent performance. However, it would be more convincing if ablations with compute time on the x axis were provided to compare performance between approaches. \n- The state regularization parameter epsilon is concerning to me for two reasons. 1) In some MDPs, taking a linear combination of states may not be sensible. For example, in a maze, two states may be feasible but their linear combination may be inside a wall. Alternatively, if the state is an image, linear combinations do not make sense. 2) The ablation of epsilon raises many questions. The ablation makes it seem like a small epsilon around 0.1 is best, with degrading performance is epsilon = 0 or 1. Furthermore, the algorithm is very sensitive to this choice. I would like to see a more thorough ablation on epsilon at a higher granularity to see what exactly is going on here. Perhaps there is some quadratic relationship, but there is not enough information at the moment to conclude this in my opinion.\n- T-sne is notoriously fickle and occasionally misleading. I would be curious to see other visualization techniques to see if the expected relationships in the latent space are still preserved. For example, you could set the latent dimension to 2 so that it can be directly plotted, and see if anything interesting happens. This may degrade performance but is usually a good smell test. This is not a major issue, but is worth considering.", "other_comments_or_suggestions": "Minor:\n- Extra comma on line 84, right side. \n- Typo line 104\n- Comma line 99 right side\n- Comma 246\n- There may be some typos in appendix d1 relating to phi - and phi ~. Perhaps use a different symbol then phi ~ since phi - can be interpreted as phi with no gradients blocked.", "questions_for_authors": "1) In figure 6, it seems like many of the baselines have fixed performance, such as LDM, Varibad, RL^2. Why is this? \n2) Does the extra decoder in section D.1, which has a one hot encoding as input, cause some scaling issues? Please discuss how well this scales since one hot encodings are not generalizable, and whether you need a one hot for only training tasks or also test tasks. \n\nCongrats on a good paper. I am willing to increase my score if questions and weaknesses are sufficiently addressed.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces Task-Aware Virtual Training (TAVT), a meta-reinforcement learning approach designed to improve generalization to out-of-distribution (OOD) tasks through a combination of metric-based task representation, virtual task generation, and sample regularization. The proposed method is positioned as an extension of prior latent-context-based methods such as PEARL, with additional components aimed at improving robustness and representation quality.

Reviewers appreciated the clarity of the paper, the extensive empirical evaluation, and the ablations that demonstrate the value of each proposed component. While some viewed the overall methodological contribution as incremental, the paper was nonetheless seen as a solid step forward in the context of OOD generalization in meta-RL.

The rebuttal effectively addressed key concerns, particularly around baseline comparisons and latent representation structure, resulting in multiple score increases post-rebuttal. Remaining reservations mostly concern the computational cost and scalability of the method, as well as its current reliance on low-dimensional task distributions.

Overall, the paper makes a meaningful contribution to the field and is suitable for inclusion at ICML.
Decision: Accept (poster)