Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning

TL;DR

This work tackles the context distribution shift in offline meta-reinforcement learning by decoupling task representations from the behavior policy. It introduces Entropy Regularized Task Representation Learning (er-trl), which uses a GAN to approximate the entropy of a meta-behavior policy conditioned on task representations, thereby approximately minimizing the mutual information between z and the data-collection policy. The method combines MI regularization with distance-metric learning and BRAC-style offline RL to produce robust, task-faithful representations that improve adaptation to unseen tasks in MuJoCo benchmarks. Empirical results show improved in-distribution and out-of-distribution performance, along with enhanced latent-task interpretability and representation quality, highlighting practical gains for offline meta-RL applications.

Abstract

Offline meta-reinforcement learning aims to equip agents with the ability to rapidly adapt to new tasks by training on data from a set of different tasks. Context-based approaches utilize a history of state-action-reward transitions -- referred to as the context -- to infer representations of the current task, and then condition the agent, i.e., the policy and value function, on the task representations. Intuitively, the better the task representations capture the underlying tasks, the better the agent can generalize to new tasks. Unfortunately, context-based approaches suffer from distribution mismatch, as the context in the offline data does not match the context at test time, limiting their ability to generalize to the test tasks. This leads to the task representations overfitting to the offline training data. Intuitively, the task representations should be independent of the behavior policy used to collect the offline data. To address this issue, we approximately minimize the mutual information between the distribution over the task representations and behavior policy by maximizing the entropy of behavior policy conditioned on the task representations. We validate our approach in MuJoCo environments, showing that compared to baselines, our task representations more faithfully represent the underlying tasks, leading to outperforming prior methods in both in-distribution and out-of-distribution tasks.

Figures (12)

• Figure 1: Correlation between the distance between expert policies and performance improvement by utilizing mutual information objective. For ID tasks, the Pearson coefficient is $\rho=0.75$ with $p=0.058$, and for OOD tasks $\rho=0.69$ with $p=0.086$.
• Figure 2: Distinguishable task representation learning Visualizing task representations with t-SNE projection for Ant-Dir environment. The labels illustrate the direction in which the ant robot should move in degrees.
• Figure 3: Overview of er-trl. Solid black line illustrates forward computations and dashed lines illustrate the gradient paths in the computational graph. The encoder is trained to infer the task from context based on distance metrics and mutual information objectives.
• Figure 4: t-SNE visualization for reward changing environments.
• Figure 5: t-SNE visualization for dynamic changing environment.