DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with
  Prototypical Representations
  Top-performing Model-Based Reinforcement Learning (MBRL) agents, such as
Dreamer, learn the world model by reconstructing the image observations. Hence,
they often fail to discard task-irrelevant details and struggle to handle
visual distractions. To address this issue, previous work has proposed to
contrastively learn the world model, but the performance tends to be inferior
in the absence of distractions. In this paper, we seek to enhance robustness to
distractions for MBRL agents. Specifically, we consider incorporating
prototypical representations, which have yielded more accurate and robust
results than contrastive approaches in computer vision. However, it remains
elusive how prototypical representations can benefit temporal dynamics learning
in MBRL, since they treat each image independently without capturing temporal
structures. To this end, we propose to learn the prototypes from the recurrent
states of the world model, thereby distilling temporal structures from past
observations and actions into the prototypes. The resulting model, DreamerPro,
successfully combines Dreamer with prototypes, making large performance gains
on the DeepMind Control suite both in the standard setting and when there are
complex background distractions. Code available at
https://github.com/fdeng18/dreamer-pro .
