Prototypical Context-aware Dynamics Generalization for High-dimensional Model-based Reinforcement Learning

Junjie Wang, Yao Mu, Dong Li, qichao Zhang, Dongbin Zhao, Yuzheng Zhuang, Ping Luo, Bin Wang, Jianye HAO

TL;DR: We present a prototypical Context-aware Dynamics (ProtoCAD) model to capture the local dynamics by time consistent latent context.

Abstract: The ability to generalize different dynamics is crucial for decision-making in autonomous driving that relies on high-dimensional inputs. The latent world model provides a promising way to learn policies in a compact latent space for tasks with high-dimensional observations, however, its generalization across diverse environments with unseen dynamics remains challenging. Although the recurrent structure utilized in current advances helps to capture local dynamics, modeling only state transitions without an explicit understanding of environmental context limits the generalization ability of the dynamics model. To address this issue, we propose a Prototypical Context-Aware Dynamics (ProtoCAD) model, which captures the local dynamics by temporally consistent latent context and enables dynamics generalization in high-dimensional control tasks. ProtoCAD extracts useful contextual information with the prototypes clustered over the batch, and it benefits model-based reinforcement learning in two ways: 1)~A temporally consistent prototypes regularizer is utilized, which encourages the prototype assignments produced for different temporal parts of the same latent trajectory to be temporally consistent instead of comparing the features; 2)~A context representation is designed, which combines both projection embedding of latent states and aggregated prototypes and can significantly improve the dynamics generalization ability. Extensive experiments show that ProtoCAD surpasses existing methods in terms of dynamics generalization.

Track: Original Contribution

Type: PDF