Learning Task-Aware Abstract Representations for Meta-Reinforcement Learning

Louk van Remmerden, Zhao Yang, Shujian Yu, Mark Hoogendoorn, Vincent Francois-Lavet

Abstract: A central challenge in meta-reinforcement learning (meta-RL) is enabling agents trained on a set of environments to generalize to new, related tasks without requiring full policy retraining. Existing model-free approaches often rely on context-conditioned policies learned via encoder networks. However, these context encoders are prone to overfitting to the training environments, resulting in poor out-of-sample performance on unseen tasks. To address this issue, we adopt an alternative approach that uses an abstract representation model to learn augmented, task-aware abstract states. We achieve this by introducing a novel architecture that offers greater flexibility than existing recurrent network-based approaches. In addition, we optimize our model with multiple loss terms that encourage predictive, task-aware representations in the abstract state space. Our method simplifies the learning problem and provides a flexible framework that can be readily combined with any off-the-shelf reinforcement learning algorithm. We provide theoretical guarantees alongside empirical results, showing strong generalization performance across classical control and robotic meta-RL benchmarks, on par with state-of-the-art meta-RL methods and significantly better than non-meta RL approaches.

Submission Length: Regular submission (no more than 12 pages of main content)

Previous TMLR Submission Url: https://openreview.net/forum?id=mkzZ0ndscN&referrer=%5BAuthor%20Console%5D(%2Fgroup%3Fid%3DTMLR%2FAuthors%23your-submissions)

Changes Since Last Submission:

Hi Laurent,

Apologies for the oversight, we used "M" and "L" instead of the direct references. In this version of the manuscript, all appendix references should now be clickable.

Please let us know if there's anything else we can improve.

All the best,

The authors

Code: https://github.com/ljsmalbil/EMERALD

Assigned Action Editor: ~Laurent_Charlin1

Submission Number: 5644