# Disentangling Policy from Offline Task Representation Learning via Adversarial Data Augmentation

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Decision transformers and sequence modeling, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 28 of its 43 references that are in the map.
- First posted on arXiv 12 Mar 2024.

## Placed in

- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [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)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [Residual Reinforcement Learning for Robot Control](https://mlsanity.com/paper/1812.03201/) (2018)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling](https://mlsanity.com/paper/2209.14548/) (2022)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Offline Meta-Reinforcement Learning with Advantage Weighting](https://mlsanity.com/paper/2008.06043/) (2020)
- [RL-CycleGAN: Reinforcement Learning Aware Simulation-To-Real](https://mlsanity.com/paper/2006.09001/) (2020)
- [MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL](https://mlsanity.com/paper/2305.19923/) (2023)
- [Conservative Data Sharing for Multi-Task Offline Reinforcement Learning](https://mlsanity.com/paper/2109.08128/) (2021)
- [Offline Meta-Reinforcement Learning with Online Self-Supervision](https://mlsanity.com/paper/2107.03974/) (2021)
- [Model-based Reinforcement Learning: A Survey](https://mlsanity.com/paper/2006.16712/) (2020)
- [Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2206.10442/) (2022)
- [Flow to Control: Offline Reinforcement Learning with Lossless Primitive Discovery](https://mlsanity.com/paper/2212.01105/) (2022)

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

- [Energy-Guided Diffusion Sampling for Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2407.12448/) (2024)
- [Model-Based Offline Reinforcement Learning with Adversarial Data Augmentation](https://mlsanity.com/paper/2503.20285/) (2025)

## Abstract (arXiv)

Offline meta-reinforcement learning (OMRL) proficiently allows an agent to tackle novel tasks while solely relying on a static dataset. For precise and efficient task identification, existing OMRL research suggests learning separate task representations that be incorporated with policy input, thus forming a context-based meta-policy. A major approach to train task representations is to adopt contrastive learning using multi-task offline data. The dataset typically encompasses interactions from various policies (i.e., the behavior policies), thus providing a plethora of contextual information regarding different tasks. Nonetheless, amassing data from a substantial number of policies is not only impractical but also often unattainable in realistic settings. Instead, we resort to a more constrained yet practical scenario, where multi-task data collection occurs with a limited number of policies. We observed that learned task representations from previous OMRL methods tend to correlate spuriously with the behavior policy instead of reflecting the essential characteristics of the task, resulting in unfavorable out-of-distribution generalization. To alleviate this issue, we introduce a novel algorithm to disentangle the impact of behavior policy from task representation learning through a process called adversarial data augmentation. Specifically, the objective of adversarial data augmentation is not merely to generate data analogous to offline data distribution; instead, it aims to create adversarial examples designed to confound learned task representations and lead to incorrect task identification. Our experiments show that learning from such adversarial samples significantly enhances the robustness and effectiveness of the task identification process and realizes satisfactory out-of-distribution generalization.