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

TL;DR

This paper introduces a context-aware dynamics modeling approach for model-based reinforcement learning that improves generalization across different environments by encoding local dynamics into a latent vector.

Contribution

The paper proposes a two-stage learning framework with a novel loss function to encode environment-specific dynamics, enhancing generalization in model-based RL.

Findings

Achieves superior generalization in robotics and control tasks

Outperforms existing RL schemes in diverse environments

Effective encoding of local dynamics into latent vectors

Abstract

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.