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

  • 2024-09-30 10:24:30
  • Hai Zhang, Boyuan Zheng, Tianying Ji, Jinhang Liu, Anqi Guo, Junqiao Zhao, Lanqing Li
  • 0

Abstract

Offline meta reinforcement learning (OMRL) has emerged as a promisingapproach for interaction avoidance and strong generalization performance byleveraging pre-collected data and meta-learning techniques. Previouscontext-based approaches predominantly rely on the intuition that alternatingoptimization between the context encoder and the policy can lead to performanceimprovements, as long as the context encoder follows the principle ofmaximizing the mutual information between the task and the task representation($I(Z;M)$) while the policy adopts the standard offline reinforcement learning(RL) algorithms conditioning on the learned task representation. Despitepromising results, the theoretical justification of performance improvementsfor such intuition remains underexplored. Inspired by the return discrepancyscheme in the model-based RL field, we find that the previous optimizationframework can be linked with the general RL objective of maximizing theexpected return, thereby providing a feasible explanation concerningperformance improvements. Furthermore, after scrutinizing this optimizationframework, we find it ignores the impacts stemming from the variation of thetask representation in the alternating optimization process, which may lead toperformance improvement collapse. We name this issue \underline{taskrepresentation shift} and theoretically prove that the monotonic performanceimprovements can be guaranteed with appropriate context encoder updates. We setdifferent manners to rein in the task representation shift on three widelyadopted training objectives concerning maximizing $I(Z;M)$ across differentdata qualities. Empirical results show that reining in the task representationshift can indeed improve performance. Our work opens up a new avenue for OMRL,leading to a better understanding between the performance and the taskrepresentation.

 

Quick Read (beta)

loading the full paper ...