Prototypical Context-Aware Dynamics for Generalization in Visual Control With Model-Based Reinforcement Learning

Abstract

The latent world model, which efficiently represents high-dimensional observations within a latent space, has shown promise in reinforcement learning-based policies for visual control tasks. Due to a lack of clear environmental context comprehension, its applicability in a variety of contexts with unknown dynamics is constrained. We propose a prototypical context- aware dynamics (ProtoCAD) model to address this issue. This model captures local dynamics using temporally consistent latent contexts and aids generalization in visual control tasks. By grouping prototypes over historical experiences, ProtoCAD collects useful contextual information that improves model-based reinforcement learning dynamics generalization in two ways. First, to guarantee the consistency of prototype assignments for various temporal segments of the same latent trajectory, a temporally consistent prototypes regularizer is used. Then, a context representation is devised to combine the aggregated prototype with the projection embedding of latent states. According to extensive trials, ProtoCAD outperforms competing approaches in terms of dynamics generalization for visual robotic control and autonomous driving applications.