# Scrutinize What We Ignore: Reining In Task Representation Shift Of Context-Based Offline Meta Reinforcement Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 39 of its 46 references that are in the map.
- First posted on arXiv 20 May 2024.

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

- [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)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [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)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [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)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Algorithmic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees](https://mlsanity.com/paper/1807.03858/) (2018)
- [Model-based Reinforcement Learning and the Eluder Dimension](https://mlsanity.com/paper/1406.1853/) (2014)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [Offline Meta-Reinforcement Learning with Advantage Weighting](https://mlsanity.com/paper/2008.06043/) (2020)
- [Online and Offline Reinforcement Learning by Planning with a Learned Model](https://mlsanity.com/paper/2104.06294/) (2021)
- [Efficient Model-Based Reinforcement Learning through Optimistic Policy Search and Planning](https://mlsanity.com/paper/2006.08684/) (2020)
- [S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learning](https://mlsanity.com/paper/2103.06326/) (2021)
- [A Workflow for Offline Model-Free Robotic Reinforcement Learning](https://mlsanity.com/paper/2109.10813/) (2021)
- [Leveraging Factored Action Spaces for Efficient Offline Reinforcement Learning in Healthcare](https://mlsanity.com/paper/2305.01738/) (2023)
- [Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2206.10442/) (2022)
- [Provably Efficient Model-Free Constrained RL with Linear Function Approximation](https://mlsanity.com/paper/2206.11889/) (2022)
- [Efficient Model-free Reinforcement Learning in Metric Spaces](https://mlsanity.com/paper/1905.00475/) (2019)

## Abstract (arXiv)

Offline meta reinforcement learning (OMRL) has emerged as a promising approach for interaction avoidance and strong generalization performance by leveraging pre-collected data and meta-learning techniques. Previous context-based approaches predominantly rely on the intuition that alternating optimization between the context encoder and the policy can lead to performance improvements, as long as the context encoder follows the principle of maximizing the mutual information between the task variable M and its latent representation Z (I(Z;M)) while the policy adopts the standard offline reinforcement learning (RL) algorithms conditioning on the learned task representation.Despite promising results, the theoretical justification of performance improvements for such intuition remains underexplored.Inspired by the return discrepancy scheme in the model-based RL field, we find that the previous optimization framework can be linked with the general RL objective of maximizing the expected return, thereby explaining performance improvements. Furthermore, after scrutinizing this optimization framework, we observe that the condition for monotonic performance improvements does not consider the variation of the task representation. When these variations are considered, the previously established condition may no longer be sufficient to ensure monotonicity, thereby impairing the optimization process.We name this issue task representation shift and theoretically prove that the monotonic performance improvements can be guaranteed with appropriate context encoder updates.Our work opens up a new avenue for OMRL, leading to a better understanding between the task representation and performance improvements.