Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning
This addresses the problem of generalization in model-based RL for robotics and control, offering an incremental improvement over existing methods.
The paper tackles the challenge of learning a global dynamics model that generalizes across different environments in model-based reinforcement learning by decomposing the task into learning a context latent vector and predicting next states, achieving superior generalization in simulated robotics and control tasks.
Model-based reinforcement learning (RL) enjoys several benefits, such as data-efficiency and planning, by learning a model of the environment's dynamics. However, learning a global model that can generalize across different dynamics is a challenging task. To tackle this problem, we decompose the task of learning a global dynamics model into two stages: (a) learning a context latent vector that captures the local dynamics, then (b) predicting the next state conditioned on it. In order to encode dynamics-specific information into the context latent vector, we introduce a novel loss function that encourages the context latent vector to be useful for predicting both forward and backward dynamics. The proposed method achieves superior generalization ability across various simulated robotics and control tasks, compared to existing RL schemes.