Learning Task-Aware Abstract Representations for Meta-Reinforcement Learning

TMLR Paper5644 Authors

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 on 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 more 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 easily 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.

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:

Dear Editors,

We received a desk-reject (Submission Number: 5583) due to an error on our part with regard to the TMLR's stylefile format (the margins were not correct). We apologise for this oversight and remedied it straightaway.

Additionally, since we had some space left (due to the wider margins), we made some small stylistic changes (e.g. included more itemised lists).

Kind regards, on behalf of the co-authors,

Louk van Remmerden

Assigned Action Editor: ~Laurent_Charlin1

Submission Number: 5644