# Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks

arXiv: https://arxiv.org/abs/2502.02834 · HTML: https://mlsanity.com/paper/2502.02834/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 49 of its 66 references that are in the map.
- First posted on arXiv 5 Feb 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Diversity is All You Need: Learning Skills without a Reward Function](https://mlsanity.com/paper/1802.06070/) (2018)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Meta Learning Shared Hierarchies](https://mlsanity.com/paper/1710.09767/) (2017)
- [Few-Shot Preference Learning for Human-in-the-Loop RL](https://mlsanity.com/paper/2212.03363/) (2022)
- [Meta-Q-Learning](https://mlsanity.com/paper/1910.00125/) (2019)
- [Continuous Meta-Learning without Tasks](https://mlsanity.com/paper/1912.08866/) (2019)
- [Transformers are Meta-Reinforcement Learners](https://mlsanity.com/paper/2206.06614/) (2022)
- [AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents](https://mlsanity.com/paper/2310.09971/) (2023)
- [Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2204.13060/) (2022)
- [Skill-based Meta-Reinforcement Learning](https://mlsanity.com/paper/2204.11828/) (2022)
- [Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2206.10442/) (2022)
- [Model-based Adversarial Meta-Reinforcement Learning](https://mlsanity.com/paper/2006.08875/) (2020)
- [Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning](https://mlsanity.com/paper/2009.13891/) (2020)
- [Provable Benefit of Multitask Representation Learning in Reinforcement Learning](https://mlsanity.com/paper/2206.05900/) (2022)
- [Meta-Reinforcement Learning Robust to Distributional Shift via Model Identification and Experience Relabeling](https://mlsanity.com/paper/2006.07178/) (2020)
- [Robust Imitation Learning against Variations in Environment Dynamics](https://mlsanity.com/paper/2206.09314/) (2022)
- [Block Contextual MDPs for Continual Learning](https://mlsanity.com/paper/2110.06972/) (2021)

## Built on by (most cited first, 3 of 3)

- [Self-Improving Skill Learning for Robust Skill-based Meta-Reinforcement Learning](https://mlsanity.com/paper/2502.03752/) (2025)
- [An Information-Theoretic Analysis of OOD Generalization in Meta-Reinforcement Learning](https://mlsanity.com/paper/2510.23448/) (2025)
- [Bridging Domain Gaps with Target-Aligned Generation for Offline Reinforcement Learning](https://mlsanity.com/paper/2605.13054/) (2026)

## Abstract (arXiv)

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.