Policy Contrastive Decoding for Robotic Foundation Models

TL;DR

This work tackles the generalization gap of robotic foundation models by identifying reliance on spurious correlations in pre-training data. It introduces Policy Contrastive Decoding (PCD), a training-free, plug-in method that reinforces object-focused cues by contrasting action distributions from original versus object-masked observations, supported by Track2Mask for automatic masking and KDE-based Probabilistic Modeling (KDE-PM) for diffusion-based policies. Empirically, PCD yields substantial gains across 15 tasks and three policies in simulation (up to 50.6% for OpenVLA, 29.7% for Octo, 8.9% for $\pi_0$) and real-world (108% average improvement) settings, with manageable but non-negligible inference-time overhead. The approach demonstrates strong generalization benefits and practical applicability, while highlighting avenues for efficiency improvements and future training-time mitigation of spurious correlations.

Abstract

Robotic foundation models, or generalist robot policies, hold immense potential to enable flexible, general-purpose and dexterous robotic systems. Despite their advancements, our empirical experiments reveal that existing robot policies are prone to learning spurious correlations from pre-training trajectories, adversely affecting their generalization capabilities beyond the training data. To tackle this, we propose a novel Policy Contrastive Decoding (PCD) approach, which redirects the robot policy's focus toward object-relevant visual clues by contrasting action probability distributions derived from original and object-masked visual inputs. As a training-free method, our PCD can be used as a plugin to improve different types of robot policies without needing to finetune or access model weights. We conduct extensive experiments on top of three open-source robot policies, including the autoregressive policy OpenVLA and the diffusion-based policies Octo and $π_0$. The obtained results in both simulation and real-world environments prove PCD's flexibility and effectiveness, e.g., PCD enhances the state-of-the-art policy $π_0$ by 8.9% in the simulation environment and by 108% in the real-world environment. Code and demos are publicly available at: https://Koorye.github.io/proj/PCD.

Figures (8)

• Figure 1: Robot policies tend to spuriously correlate task-irrelevant features with actions, compromising their ability to generalize to unseen scenarios. As observed, changing the light position from (a) to (b) and the drawer handle position from (a) to (c) results in 36% and 32% drops in the performance of the baseline policy OpenVLAkim2024openvla, respectively. (d) Attention map. More results are in Section \ref{['sect.factor']} and Appendix\ref{['app.spurious']}.
• Figure 2: Overview of our Proposed Policy Contrastive Decoding (PCD) approach. PCD serves as a plugin to redirect the robot policy’s focus toward object-relevant visual cues by contrasting action probability distributions derived from original observations ${p}$ and object-masked observations $\hat{p}$. For illustrative purposes, we visualize the predictions only in the $\Delta x$ and $\Delta y$ dimensions of the robot action space [$\Delta x$, $\Delta y$, $\Delta z$, rot$_x$, rot$_y$, rot$_z$, gripper].
• Figure 3: Real-world Performance. The target objects in the initial observation are automatically annotated by Grounding DINO liu2024grounding. PCD delivers a remarkable 108% performance improvement on the baseline, though it incurs a 24% increase in time cost.
• Figure 4: Ablation studies on (a) the hyperparameter $\alpha$ in Eq. (\ref{['eq.pcd']}), (b) the object detection schemes and (c) object inpainting strategies in Track2Mask. $\alpha=0$ in (a) and the black dotted lines in (b)(c) represent the performance of the baseline policies. The results are averaged over the 9 simulation tasks. PCD consistently improves the three policies when $\alpha>0$ and exhibits low sensitivity to changes in off-the-shelf object detection and inpainting strategies.
• Figure 5: Performance of baseline policies integrated w/ or w/o our proposed PCD approach in unseen testing scenarios. Dotted lines represent the original performance. OpenVLAkim2024openvla and $\pi_0$black2024pi0 are used as baselines in simulation and real-world environments, respectively. PCD effectively mitigates the side effects of various types of spurious correlations, boosting the robot policy's generalization to novel scenarios.