Pre-trained Visual Dynamics Representations for Efficient Policy Learning

TL;DR

This work proposes Pre-trained Visual Dynamics Representations (PVDR) to bridge the domain gap between videos and downstream tasks for efficient policy learning and conducts experiments to verify that PVDR is an effective form for pre-training with videos to promote policy learning.

Abstract

Pre-training for Reinforcement Learning (RL) with purely video data is a valuable yet challenging problem. Although in-the-wild videos are readily available and inhere a vast amount of prior world knowledge, the absence of action annotations and the common domain gap with downstream tasks hinder utilizing videos for RL pre-training. To address the challenge of pre-training with videos, we propose Pre-trained Visual Dynamics Representations (PVDR) to bridge the domain gap between videos and downstream tasks for efficient policy learning. By adopting video prediction as a pre-training task, we use a Transformer-based Conditional Variational Autoencoder (CVAE) to learn visual dynamics representations. The pre-trained visual dynamics representations capture the visual dynamics prior knowledge in the videos. This abstract prior knowledge can be readily adapted to downstream tasks and aligned with executable actions through online adaptation. We conduct experiments on a series of robotics visual control tasks and verify that PVDR is an effective form for pre-training with videos to promote policy learning.

Figures (18)

• Figure 1: Illustration of the pre-training video dataset and the downstream environments we used for experimental evaluations.
• Figure 2: Illustration of video prediction with dVAE. In this paradigm, a dVAE is pre-trained to compress frames from raw pixel-level space (upper block) into discrete latent space (lower block) for effective prediction.
• Figure 3: Illustration of PVDR workflow for different stages and usages. Four modules, Visual Dynamics Encoder, Visual Dynamics Decoder, Visual Dynamics Prior, and Action Alignment Module, are included in PVDR. The learning workflow is on the left, while the inference workflow is on the right. In the pre-training stage (blue block), Encoder, Decoder, and Prior are pre-trained to capture visual dynamics representations. During the online adaptation (gray block), the Action Alignment Module is integrated with the other three modules. Planning-based inference (green block) is performed to choose the plans closest to the goal and execute aligned actions.
• Figure 4: Learning curves of PVDR compared with five other baselines on 12 Meta-World tasks measured on success rate. The solid line and shaded regions represent the mean and variance of the performance across five runs with different seeds.
• Figure 5: Learning curves of PVDR's ablation studies on 4 Meta-World tasks measured on success rate. The dashed line illustrates the mean success rate of converged PVDR. The solid line and the shaded regions represent the mean and variance of performance across five runs with different seeds. PVDR w/o prior opt, PVDR w/o act sl, PVDR w/o act rl, and PVDR random refer to PVDR without goal-oriented term in $\mathcal{L}_{prior}$, PVDR without $\mathcal{L}_{act}$, PVDR without PPO loss, and PVDR with random representation.