Latent Policy Steering with Embodiment-Agnostic Pretrained World Models

TL;DR

This paper introduces Latent Policy Steering, a method that pretrains world models using embodiment-agnostic visual motion representations like optical flow, enabling improved robot visuomotor policies with limited target data.

Contribution

It proposes a novel approach to leverage diverse datasets across embodiments by pretraining world models with optical flow, then finetuning for specific tasks to enhance policy performance.

Findings

LPS improves behavior-cloned policies by 10.6% on average across four tasks.

In real-world tests, LPS achieves 70% and 44% relative improvements with 30-50 and 60-100 demonstrations respectively.

Pretraining with embodiment-agnostic data enables effective transfer and data efficiency.

Abstract

The performance of learned robot visuomotor policies is heavily dependent on the size and quality of the training dataset. Although large-scale robot and human datasets are increasingly available, embodiment gaps and mismatched action spaces make them difficult to leverage. Our main insight is that skills performed across different embodiments produce visual similarities in motions that can be captured using off-the-shelf action representations such as optical flow. Moreover, World Models (WMs) can leverage sub-optimal data since they focus on modeling dynamics. In this work, we aim to improve visuomotor policies in low-data regimes by first pretraining a WM using optical flow as an embodiment-agnostic action representation to leverage accessible or easily collected data from multiple embodiments (robots, humans). Given a small set of demonstrations on a target embodiment, we finetune the WM on this data to better align the WM predictions, train a base policy, and learn a robust value function. Using our finetuned WM and value function, our approach evaluates action candidates from the base policy and selects the best one to improve performance. Our approach, which we term Latent Policy Steering (LPS), improves behavior-cloned policies by 10.6% on average across four Robomimic tasks, even though most of the pretraining data comes from the real world. In the real-world experiments, LPS achieves larger gains: 70% relative improvement with 30-50 target-embodiment demonstrations, and 44% relative improvement with 60-100 demonstrations, compared to a behavior-cloned baseline. Qualitative results can be found on the website: https://yiqiwang8177.github.io/LatentPolicySteering/.