Technical Summary: TapSampling

Problem Statement

Existing research in embodied control has largely focused on scaling training data and model size to improve robotic manipulation policies. While non-deterministic generative models (e.g., diffusion models and autoregressive Large Language Models) have demonstrated remarkable capabilities, their performance is fundamentally constrained by the single-shot inference paradigm. Due to the stochastic nature of these models, identical environment conditions can yield successful outcomes in some instances and failures in others. Current inference-time strategies for robotics often suffer from two main limitations:

1. Sampling Inefficiency or Fidelity: Repeatedly sampling directly from policy models significantly increases inference latency (by an order of magnitude), while sampling from simple Gaussian distributions ignores inter-dimensional correlations in action chunks, leading to candidates that deviate from the true action distribution.
2. Verification Challenges: There is a lack of off-the-shelf models capable of understanding and interpreting low-level robotic actions. Existing verification methods often rely on hand-crafted functions, offline reinforcement learning, or preference learning, which may lack interpretability or are tightly coupled to specific policy architectures, limiting their plug-and-play applicability.

Methodology

The authors propose TapSampling, a plug-and-play, policy-agnostic inference-time sampling framework designed to improve generalist robotic policies without further fine-tuning. The framework operates through two primary components:

1. Efficient Action Sampling with Learned Posterior

To address the trade-off between sampling efficiency and action quality, TapSampling introduces an Action-VAE (Variational Autoencoder).

• Architecture: The Action-VAE utilizes a transformer-based encoder and decoder.
• Process:
  • A small set of initial actions ($N$) is sampled from the base policy $\pi$.
  • The encoder compresses these actions into a latent Gaussian distribution (mean $\mu$ and diagonal covariance $\sigma^2$).
  • A mixed posterior distribution is constructed by averaging the latent distributions of the initial samples.
  • An arbitrary number of new latent samples are drawn from this compressed posterior and decoded into candidate actions.
• Benefit: This approach captures intra-chunk correlations among action dimensions, ensuring generated candidates remain close to the true policy distribution while being significantly faster than repeated policy sampling (approx. 5$\times$ faster).

2. Interpretable Action Selection via Task-Progress Verification

To select the optimal action among candidates, TapSampling employs a Task-Progress-Understanding Verifier.

• Training Data Construction: The verifier is trained using the intrinsic sequential structure of expert robotic trajectories. By assuming task progress grows linearly over time, the authors construct positive samples (forward action sequences leading to progress) and negative samples (reversed action sequences leading to regression) directly from existing datasets, requiring no additional manual annotation.
• Architecture: The verifier is built upon a Vision-Language-Action (VLA) backbone (specifically Qwen2.5-0.5B). It takes the current state, task instruction, and a query action as input to predict the expected change in task progress ($\Delta p$).
• Selection Mechanism: The predicted score carries explicit semantic meaning (e.g., a negative score indicates hindrance, a large positive score indicates rapid progress). During inference, candidates with scores below a threshold are discarded, and a score-weighted average of the remaining candidates is computed to produce the final action.
• Efficiency: The VLM backbone is executed only once per observation-instruction pair. Hidden states are shared across all candidates, which are then evaluated in parallel through a lightweight action head, enabling high-throughput verification.

Key Contributions

1. Policy-Agnostic Framework: A plug-and-play system that enhances pretrained generalist policies without requiring additional fine-tuning.
2. Learned Posterior Sampling: An Action-VAE-based strategy that encodes actions into a compressed distribution, allowing for the efficient generation of diverse, high-quality action candidates that adhere to the policy's distribution.
3. Task-Progress Verification: A verifier trained to predict the impact of low-level actions on task progress, enabling interpretable action selection based on explicit semantic scores rather than opaque reward signals.
4. Empirical Validation: Demonstration of substantial performance improvements across diverse simulated and real-world environments.

Experimental Results

The authors evaluated TapSampling on both simulated benchmarks (CALVIN, LIBERO) and real-world robotic setups.

• CALVIN Benchmark (Simulated): TapSampling improved the average success length for multiple non-deterministic policies:
  • Diffusion Policy: Increased from 2.41 to 2.58.
  • OpenVLA: Increased from 3.30 to 3.51.
  • VPP: Increased from 4.39 to 4.46.
  • Improvements were consistent across different task steps (1–5), indicating better long-horizon stability.
• LIBERO-Long Benchmark (Simulated): Applied to the $\pi0.5$ policy, the method improved the success rate from 96.8% to 98.0%.
• Real-World Experiments: Using a Franka Research 3 arm, TapSampling improved the average success rate of the $\pi0$ policy from 78.3% to 83.3%. Notably, it enhanced precision in tasks requiring spatial alignment and contact-aware placement (e.g., stacking and placing objects).
• Efficiency Analysis:
  • Sampling: Learned Posterior Sampling incurred negligible additional latency (<0.01s) compared to Gaussian sampling, while being significantly faster than repeated Policy Sampling.
  • Verification: TapSampling achieved approximately a 12$\times$ speedup in verification latency compared to RoboMonkey (which uses LLaVA-7B) when evaluating 16 candidates, due to its batched evaluation architecture.
  • Fidelity: Maximum Mean Discrepancy (MMD) analysis confirmed that actions generated via Learned Posterior Sampling align more closely with the true policy distribution than those from Gaussian sampling.

Significance and Claims

The paper posits that allocating computation to inference time offers an effective, model-agnostic pathway to improve the reliability and performance of generalist robotic policies, serving as a complementary axis to scaling training data.

The authors claim that TapSampling addresses the instability inherent in single-shot inference by:

1. Providing a mechanism to correct failures caused by stochastic variance without retraining the policy.
2. Enabling interpretable action selection through a verifier that understands task progress, rather than relying on black-box scoring.
3. Balancing efficiency and effectiveness through the Learned Posterior Sampling strategy, which avoids the high latency of repeated policy queries while maintaining high action fidelity.

The work suggests that inference-time strategies, particularly those grounded in task-progress understanding, are a viable and necessary direction for deploying robust robotic manipulation systems in complex, real-world scenarios.