Policy Contrastive Decoding for Robotic Foundation Models

1. Problem Statement

Robotic foundation models (generalist robot policies) have shown great potential for flexible, general-purpose manipulation. However, the authors identify a critical flaw: these models often learn spurious correlations from pre-training trajectories. Instead of focusing on task-relevant object features, policies rely on irrelevant background cues (e.g., lighting, textures, or specific background objects) to predict actions.

• Consequence: When the deployment environment shifts (e.g., background changes, lighting adjustments, or object repositioning), the policy's performance degrades significantly because the learned "shortcuts" are no longer valid.
• Current Limitations: Existing solutions often require retraining or fine-tuning, which is computationally expensive and requires access to model weights. There is a lack of training-free, plug-and-play methods to rectify these inference-time failures.

2. Methodology: Policy Contrastive Decoding (PCD)

The authors propose Policy Contrastive Decoding (PCD), a training-free inference-time plugin designed to redirect a robot policy's attention from spurious features to object-relevant visual clues.

Core Mechanism

PCD operates by contrasting the action probability distributions generated from two different visual inputs:

1. Original Observation ($o_i$): The standard visual input.
2. Object-Masked Observation ($\hat{o}_i$): The same visual input where the target task-relevant objects are masked (removed/inpainted).

The method computes a new, contrastive action probability distribution ($\pi^*$) using the following formula:
$$\pi^*_\theta(a_i | \ell, o_i) = \frac{1}{C} \cdot \pi_\theta(a_i | \ell, o_i) \cdot \left( \frac{\pi_\theta(a_i | \ell, o_i)}{\pi_\theta(a_i | \ell, \hat{o}_i)} \right)^\alpha$$

• Logic: If the policy relies on spurious features (background), the probability distribution will be similar for both $o_i$ and $\hat{o}_i$, resulting in a ratio near 1 (no change). If the policy relies on the object, the distribution for $\hat{o}_i$ will differ significantly, amplifying the correct action probability in the final output.
• Hyperparameter ($\alpha$): Controls the strength of the contrast. $\alpha=0$ recovers the baseline policy.

Key Technical Components

To make PCD applicable to diverse robot policies, the authors introduce two specific strategies:

1. Tracking-to-Mask (Track2Mask):

   • Goal: Automatically generate object-masked observations for every frame in a trajectory without manual intervention.
   • Process:
     • Annotation: Target objects in the initial frame are identified using human prompts (Point/Box) or off-the-shelf open-vocabulary detectors (e.g., Grounding DINO).
     • Tracking: The SAM2 (Segment Anything Model 2) tracks and segments these objects across subsequent frames.
     • Inpainting: The segmented objects are inpainted (removed) to create the masked observation $\hat{o}_i$.

2. KDE-based Probabilistic Modeling (KDE-PM):

   • Challenge: Autoregressive policies (e.g., OpenVLA) output token probabilities directly, but diffusion-based policies (e.g., Octo, $\pi_0$) generate actions via a denoising process and do not natively output a probability distribution $\pi_\theta(a|o)$.
   • Solution: The authors use Kernel Density Estimation (KDE) to approximate the action probability distribution. They sample $N$ candidate actions from the diffusion policy and fit a Gaussian kernel density to these samples to estimate $\pi_\theta(a|o)$ and $\pi_\theta(a|\hat{o})$, enabling the contrastive decoding formula to work with diffusion models.

3. Key Contributions

1. First Training-Free Solution: PCD is the first work to propose a plug-and-play, training-free approach to mitigate spurious correlations in robot policies during inference.
2. Model Agnostic: It functions as a black-box plugin compatible with both autoregressive (OpenVLA) and diffusion-based (Octo, $\pi_0$) policies without requiring access to pre-trained weights or fine-tuning.
3. Robustness to Distribution Shift: By explicitly contrasting object-relevant vs. object-irrelevant inputs, PCD forces the policy to focus on causal features, significantly improving generalization to unseen scenarios.

4. Experimental Results

The authors evaluated PCD on 15 diverse tasks across simulation (SIMPLER) and real-world environments using three state-of-the-art policies: OpenVLA, Octo, and $\pi_0$.

• Simulation Performance:
  • OpenVLA: Improved success rate by 50.6% on average.
  • Octo: Improved success rate by 29.7% on average.
  • $\pi_0$: Improved success rate by 8.9% on average (despite $\pi_0$ already being a strong baseline).
• Real-World Performance:
  • On real-world manipulation tasks using the $\pi_0$ policy, PCD achieved a massive 108% improvement in success rate compared to the baseline.
  • The method proved robust against various distractors (lighting changes, background clutter, object position shifts).
• Ablation Studies:
  • The method is robust to different object detection strategies (Human prompts vs. Grounding DINO vs. YOLO World).
  • It is robust to different inpainting algorithms (LaMa performed best).
  • The hyperparameter $\alpha$ shows a clear positive impact when $>0$.

5. Significance and Limitations

• Significance: PCD addresses a fundamental weakness in current robotic foundation models (reliance on spurious correlations) without the heavy cost of retraining. It demonstrates that simple inference-time interventions can drastically improve generalization and safety in real-world deployment.
• Limitations:
  • Computational Overhead: PCD requires a second forward pass (for the masked image) and object tracking/inpainting, increasing inference latency by approximately 24-100% (depending on the policy and parallelization). However, the authors argue this is acceptable given the massive gains in success rates.
  • Scope: Currently focuses on inference-time correction; future work aims to address spurious correlations during the training phase.

Conclusion: The paper presents a highly effective, versatile framework (PCD) that significantly enhances the reliability and generalization of robotic foundation models by forcing them to attend to the correct visual features, bridging the gap between simulation performance and real-world robustness.