# On Context Distribution Shift in Task Representation Learning for Offline Meta RL

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 34 references that are in the map.
- First posted on arXiv 1 Apr 2023.

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

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [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)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere](https://mlsanity.com/paper/2005.10242/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Meta-Learning in Neural Networks: A Survey](https://mlsanity.com/paper/2004.05439/) (2020)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Contrastive Learning with Hard Negative Samples](https://mlsanity.com/paper/2010.04592/) (2020)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Understanding the Behaviour of Contrastive Loss](https://mlsanity.com/paper/2012.09740/) (2020)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [A Survey on Offline Reinforcement Learning: Taxonomy, Review, and Open Problems](https://mlsanity.com/paper/2203.01387/) (2022)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [SMARTS: Scalable Multi-Agent Reinforcement Learning Training School for Autonomous Driving](https://mlsanity.com/paper/2010.09776/) (2020)
- [Towards the Generalization of Contrastive Self-Supervised Learning](https://mlsanity.com/paper/2111.00743/) (2021)
- [Offline Meta-Reinforcement Learning with Advantage Weighting](https://mlsanity.com/paper/2008.06043/) (2020)
- [Meta reinforcement learning as task inference](https://mlsanity.com/paper/1905.06424/) (2019)
- [Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2206.10442/) (2022)

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

- [Efficient and Generalized end-to-end Autonomous Driving System with Latent Deep Reinforcement Learning and Demonstrations](https://mlsanity.com/paper/2401.11792/) (2024)
- [Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning](https://mlsanity.com/paper/2412.14834/) (2024)
- [Learning Top-k Subtask Planning Tree based on Discriminative Representation Pre-training for Decision Making](https://mlsanity.com/paper/2312.11027/) (2023)
- [Exploring the Link Between Bayesian Inference and Embodied Intelligence: Toward Open Physical-World Embodied AI Systems](https://mlsanity.com/paper/2507.21589/) (2025)

## Abstract (arXiv)

Offline Meta Reinforcement Learning (OMRL) aims to learn transferable knowledge from offline datasets to enhance the learning process for new target tasks. Context-based Reinforcement Learning (RL) adopts a context encoder to expediently adapt the agent to new tasks by inferring the task representation, and then adjusting the policy based on this inferred representation. In this work, we focus on context-based OMRL, specifically on the challenge of learning task representation for OMRL. We conduct experiments that demonstrate that the context encoder trained on offline datasets might encounter distribution shift between the contexts used for training and testing. To overcome this problem, we present a hard-sampling-based strategy to train a robust task context encoder. Our experimental findings on diverse continuous control tasks reveal that utilizing our approach yields more robust task representations and better testing performance in terms of accumulated returns compared to baseline methods. Our code is available at https://github.com/ZJLAB-AMMI/HS-OMRL.