Context-aware Dynamics Model for Generalization in Model-Based\n Reinforcement Learning

Abstract

Model-based reinforcement learning (RL) enjoys several benefits, such as\ndata-efficiency and planning, by learning a model of the environment's\ndynamics. However, learning a global model that can generalize across different\ndynamics is a challenging task. To tackle this problem, we decompose the task\nof learning a global dynamics model into two stages: (a) learning a context\nlatent vector that captures the local dynamics, then (b) predicting the next\nstate conditioned on it. In order to encode dynamics-specific information into\nthe context latent vector, we introduce a novel loss function that encourages\nthe context latent vector to be useful for predicting both forward and backward\ndynamics. The proposed method achieves superior generalization ability across\nvarious simulated robotics and control tasks, compared to existing RL schemes.\n