Disentangled Task Representation Learning for Offline Meta Reinforcement Learning

In this paper, we aim to address the generalization problem in Offline Meta-Reinforcement Learning (OMRL) when both task objectives and environmental parameters vary simultaneously. We propose DIsentangled TAsk Representation learning (DITAR) for OMRL, which leverages the Conditional Variational Auto-Encoder framework to disentangle task representations into distinct components for task objectives and environmental parameters, thus enhancing policy generalization across diverse tasks. We further impose orthogonality constraints to prevent overlap between the representation spaces, ensuring that each space independently captures its corresponding task component. Additionally, mutual information optimization is applied to remove redundant state-action information, focusing the representation space on task-relevant features. Experiments on the multi-task MuJoCo benchmark demonstrate that DITAR significantly outperforms previous methods, particularly in complex scenarios involving simultaneous changes in both task objectives and environmental parameters.