Research map / Reinforcement learning / Visual and representation RL
Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Task-Aware Virtual Training (TAVT), a PEARL-based meta-RL method that uses a bisimulation-metric task latent, WGAN-based task-preserving virtual context generation, and state regularization to improve generalization to OOD tasks.
Problem
Context-based meta-RL policies generalize poorly to out-of-distribution test tasks, especially when state transition dynamics vary across tasks, and existing virtual-task methods produce contexts that do not faithfully capture task characteristics.
Approach
Learns task latents that preserve a bisimulation metric (reward and W2 transition distance) with an on-off latent loss, generates virtual tasks by interpolating latents with a task-preserving loss and WGAN discriminator to make virtual contexts realistic, mixes real next states into virtual next states (state regularization), and uses virtual-task latents for exploration.
Key results
- On MetaWorld ML1-OOD, TAVT reaches 0.98±0.02 on Push-OOD-Inter and 0.92±0.08 on Push-OOD-Extra, versus at most 0.98 for AMAGO on Push-OOD-Inter and 0.83 for AMAGO on Push-OOD-Extra among baselines.
- TAVT achieves 0.99±0.01 on Reach-OOD-Extra, compared with 0.82 for VariBAD and 0.81 for CCM.
- TAVT's full-training overhead is about 18% more than the PEARL baseline, per the authors.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TAVT instead of MAML | Meta-RL · Reach | success rate: 0.97±0.02 -> 0.98±0.02 | Helped |
| Use TAVT instead of Amago | Meta-RL · Reach-OOD-Inter | success rate: 0.93±0.05 -> 0.96±0.03 | Helped |
| Use TAVT instead of VariBAD | Meta-RL · Reach-OOD-Extra | success rate: 0.82±0.11 -> 0.99±0.01 | Helped |
| Use TAVT instead of RL² | Meta-RL · Push | success rate: 0.98±0.02 -> 0.98±0.03 | No clear effect |
| Use TAVT instead of Amago | Meta-RL · Push-OOD-Inter | success rate: 0.98±0.02 -> 0.98±0.02 | No clear effect |
| Use TAVT instead of Amago | Meta-RL · Push-OOD-Extra | success rate: 0.83±0.11 -> 0.92±0.08 | Helped |
| Add virtual-task training (ablation) | TAVT · Cheetah-Vel-OOD and Walker-Mass-OOD | average return | Helped |
| Add task-preserving sample generation loss L_gen (ablation) | TAVT · Cheetah-Vel-OOD and Walker-Mass-OOD | average return | Helped |
| Add on-off latent loss (ablation) | TAVT · Cheetah-Vel-OOD and Walker-Mass-OOD | average return | Helped |
| Replace reconstruction-only VT generation with TAVT generation (ablation) | TAVT · Cheetah-Vel-OOD and Walker-Mass-OOD | average return | Helped |
| Add task-preserving sample generation loss L_gen (ablation) | TAVT · OOD tasks | context difference | Helped |
| Add on-off latent loss (ablation) | TAVT · OOD tasks | context difference | Helped |
Surprising findings
- PEARL-based methods such as CCM perform very poorly on Push tasks (0.18 success), indicating that contrastive latents may fail on some tasks.
Limitations
- Many hyperparameters need tuning, though the authors say performance is not highly sensitive.
- Assumes a parametric task distribution; does not cover non-parametric task distributions such as ML10 or ML45.
- Evaluated on standard MuJoCo and MetaWorld ML1 benchmarks only.
Experimental setup
Datasets: MuJoCo Cheetah-Vel-OOD MuJoCo Ant-Dir-2 MuJoCo Ant-Dir-4 MuJoCo Ant-Goal-OOD MuJoCo Hopper-Mass-OOD MuJoCo Walker-Mass-OOD MetaWorld ML1 Reach/Push and OOD variants
Benchmarks: MuJoCo Cheetah-Vel-OOD MuJoCo Ant-Dir-2 MuJoCo Ant-Dir-4 MuJoCo Ant-Goal-OOD MuJoCo Hopper-Mass-OOD MuJoCo Walker-Mass-OOD MetaWorld ML1 Reach/Push and OOD variants
Metrics: average return average success rate Q-function estimation bias context difference
Baselines: MAML RL$^2$ VariBAD LDM PEARL CCM Amago MIER
Builds on: PEARL Latent Dynamics Mixture VariBAD Bisimulation metric WGAN-GP
Closest papers by meaning
- Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning (AAAI 2025 · similarity 0.90)
- Generalizable Task Representation Learning for Offline Meta-Reinforcement Learning with Data Limitations (AAAI 2024 · similarity 0.90)
- MetaCARD: Meta-Reinforcement Learning with Task Uncertainty Feedback via Decoupled Context-Aware Reward and Dynamics Components (AAAI 2024 · similarity 0.89)
- AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers (NeurIPS 2024 · similarity 0.89)
- Pessimistic Latent Task-aware Optimization for Robust Offline Meta-Reinforcement Learning (NeurIPS 2026 · similarity 0.88)
- Decoupling Meta-Reinforcement Learning with Gaussian Task Contexts and Skills (AAAI 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.