Trajectory-wise Multiple Choice Learning for Dynamics Generalization in\n Reinforcement Learning

Abstract

Model-based reinforcement learning (RL) has shown great potential in various\ncontrol tasks in terms of both sample-efficiency and final performance.\nHowever, learning a generalizable dynamics model robust to changes in dynamics\nremains a challenge since the target transition dynamics follow a multi-modal\ndistribution. In this paper, we present a new model-based RL algorithm, coined\ntrajectory-wise multiple choice learning, that learns a multi-headed dynamics\nmodel for dynamics generalization. The main idea is updating the most accurate\nprediction head to specialize each head in certain environments with similar\ndynamics, i.e., clustering environments. Moreover, we incorporate context\nlearning, which encodes dynamics-specific information from past experiences\ninto the context latent vector, enabling the model to perform online adaptation\nto unseen environments. Finally, to utilize the specialized prediction heads\nmore effectively, we propose an adaptive planning method, which selects the\nmost accurate prediction head over a recent experience. Our method exhibits\nsuperior zero-shot generalization performance across a variety of control\ntasks, compared to state-of-the-art RL methods. Source code and videos are\navailable at https://sites.google.com/view/trajectory-mcl.\n