Where Should Action Generation Begin? A Learnable Source Prior for Generative Robot Policies

Technical Summary: LeaP - A Learnable Source Prior for Generative Robot Policies

Problem Statement
Generative robot policies, including diffusion-based and flow-matching approaches, typically formulate action generation as conditional sampling. Despite architectural diversity, these methods share a common limitation: they initialize action generation from an observation-independent standard Gaussian distribution, $N(0, I)$. This source distribution contains no information regarding the current robot state or task context. Consequently, the generative model must expend a significant portion of its "integration budget" to transport samples from this uninformed noise source to task-relevant actions. While recent works have attempted to replace this standard Gaussian with deterministic, observation-derived sources (e.g., point estimates from proprioception or visual latents), they largely overlook the uncertainty inherent in the source distribution. The central question addressed by this work is: Where should action generation begin? specifically investigating whether the source should be deterministic or distributional, and if distributional, whether its variance should be fixed or state-adaptive.

Methodology: LeaP
The authors propose LeaP (Learnable source Prior), a framework that replaces the standard Gaussian source with a proprioception-conditioned, learnable diagonal Gaussian over action chunks. The key components are:

1. Prior Network ($g_\phi$): A lightweight two-layer MLP that maps proprioceptive state features ($z_s$) to the mean ($\mu_\phi$) and state-adaptive variance ($\sigma^2_\phi$) of a diagonal Gaussian distribution. Crucially, the prior conditions only on proprioception, leaving the full conditioning (visual + proprioceptive) to the downstream generator.
2. Stochastic Initialization: At inference, the source sample $z_0$ is drawn from this learned Gaussian $q_\phi(z_0 | z_s)$ rather than $N(0, I)$. This provides an initialization that is both observation-informed (anchored in a state-dependent neighborhood) and stochastic.
3. Plug-and-Play Design: LeaP modifies only the source distribution. The downstream generator architecture (e.g., flow-matching or diffusion-bridge), generative dynamics, and inference solver remain unchanged.
4. Joint Training Objective: The prior network and the generator are trained end-to-end using a composite loss function:
   • Flow Matching Loss ($L_{flow}$): Supervises the velocity field $v_\theta$ to transport $z_0$ to the target action.
   • Negative Log-Likelihood ($L_{NLL}$): Supervises the conditional density of the prior, ensuring the predicted mean and variance align with expert action chunks.
   • Contrastive Alignment Loss ($L_{align}$): A CLIP-style symmetric loss that enforces instance-level correspondence between source samples and actions in a learned embedding space.

Key Contributions

• Reframing the Source Distribution: The paper identifies the source distribution as an independent and reusable design axis for generative policies, distinct from the choice of generative dynamics.
• State-Adaptive Variance: Unlike previous deterministic sources or fixed-variance heuristics, LeaP jointly learns the mean and a state-adaptive variance, explicitly modeling source uncertainty.
• Modality Selection: The authors demonstrate that proprioceptive state alone is sufficient to parameterize the prior, and that incorporating visual features into the prior actually degrades performance in their setting.
• Plug-and-Play Compatibility: The method improves both flow-matching and diffusion-bridge generators without altering their internal architectures.

Experimental Results
The authors evaluated LeaP on 15 bimanual manipulation tasks in the RoboTwin 2.0 simulation environment and three real-world tasks on a Franka Research 3 robot.

• Simulation Performance: LeaP achieved an average success rate of 81.6%, outperforming four representative baselines (deterministic-source methods A2A and VITA, a no-prior counterpart, and a diffusion-bridge policy) by 6.5 to 25.5 percentage points. The largest gap (25.5 pp) was observed against the "NoPrior" baseline, which shared the exact same generator architecture but used $N(0, I)$, confirming the gain stems from the learnable prior itself.
• Efficiency: LeaP uses significantly fewer parameters (13.22M total, with the prior head accounting for only 0.21M) compared to baselines like ACT or BridgePolicy. It also exhibits faster and more stable training convergence.
• Real-World Deployment: In real-world experiments, LeaP maintained the top performance with an 80.0% average success rate, outperforming the next best baseline (A2A) by 11.7 percentage points.
• Ablation Studies:
  • Input Modality: Conditioning the prior on proprioception alone yielded the best results; fusing visual features into the prior reduced performance.
  • Parameterization: Jointly learning mean and state-adaptive variance was critical; deterministic or fixed-variance alternatives performed worse.
  • Generators: The prior consistently improved performance across different generator backbones.

Significance and Claims
The paper claims that the source distribution is a critical, yet underexplored, design dimension in generative robot policies. By shifting from an uninformed standard Gaussian to a learnable, proprioception-conditioned stochastic prior, the generator is freed from the burden of transporting samples from a distant noise source. Instead, the prior places the initial sample in a "proprioceptively plausible region," allowing the generator to focus on precise, observation-conditioned refinement. The authors assert that this approach offers a low-overhead lever to improve policy performance, convergence speed, and robustness, complementing existing choices in generative dynamics.