Latent Policy Barrier: Learning Robust Visuomotor Policies by Staying In-Distribution

TL;DR

Latent Policy Barrier (LPB) enhances visuomotor policy robustness by maintaining in-distribution states through a two-module system, improving data efficiency and reliability without extra human effort.

Contribution

LPB introduces a novel framework inspired by Control Barrier Functions that separates expert imitation from out-of-distribution recovery in visuomotor learning.

Findings

LPB improves policy robustness in simulated and real-world tasks.

LPB enhances data efficiency, reducing the need for extensive expert data.

LPB achieves reliable manipulation without additional human correction.

Abstract

Visuomotor policies trained via behavior cloning are vulnerable to covariate shift, where small deviations from expert trajectories can compound into failure. Common strategies to mitigate this issue involve expanding the training distribution through human-in-the-loop corrections or synthetic data augmentation. However, these approaches are often labor-intensive, rely on strong task assumptions, or compromise the quality of imitation. We introduce Latent Policy Barrier, a framework for robust visuomotor policy learning. Inspired by Control Barrier Functions, LPB treats the latent embeddings of expert demonstrations as an implicit barrier separating safe, in-distribution states from unsafe, out-of-distribution (OOD) ones. Our approach decouples the role of precise expert imitation and OOD recovery into two separate modules: a base diffusion policy solely on expert data, and a dynamics model trained on both expert and suboptimal policy rollout data. At inference time, the dynamics model predicts future latent states and optimizes them to stay within the expert distribution. Both simulated and real-world experiments show that LPB improves both policy robustness and data efficiency, enabling reliable manipulation from limited expert data and without additional human correction or annotation.