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

Abstract: Generative robot policies typically begin action generation from an observation-independent standard Gaussian distribution, leaving the choice of source distribution underexplored. This work asks a simple question: where should action generation begin? We propose LeaP, a Learnable source Prior that replaces the standard Gaussian with a proprioception-conditioned diagonal Gaussian over action chunks. Parameterized by a lightweight MLP, LeaP jointly predicts the mean and state-adaptive variance of the source distribution, while keeping the downstream generator architecture and inference solver unchanged. This design provides an observation-informed yet stochastic initialization, allowing the generator to focus on precise action refinement rather than transporting samples from an uninformed noise source. On 15 RoboTwin manipulation tasks, LeaP achieves an average success rate of 81.6%, outperforming four representative baselines -- including deterministic-source methods, a no-prior counterpart, and a diffusion-bridge policy -- by 6.5 to 25.5 percentage points. The same prior consistently improves both flow-matching and diffusion-bridge generators, while using fewer parameters and converging faster. The advantage carries over to real-world deployment, where LeaP attains the best performance. These results suggest that the source distribution is an independent and reusable design axis for generative robot policies, complementary to the choice of generative dynamics.

• The paper introduces LeaP, a proprioception-conditioned learnable source prior that significantly boosts generative robot policy performance.
• It replaces a static Gaussian with an adaptive diagonal Gaussian trained via negative log-likelihood and contrastive alignment, enhancing state-adaptive action generation.
• Empirical results on 15 RoboTwin tasks and real-robot deployments show improvements of up to 25.5 percentage points with minimal parameter overhead.

A Learnable Source Prior for Generative Robot Policies: The LeaP Framework

Introduction

This work systematically re-evaluates the foundational premise in generative robot policy learning regarding the choice of source distribution for action generation. While state-of-the-art diffusion and flow-matching policies typically begin from an uninformed, observation-independent standard Gaussian ($N(0, I)$), this paper demonstrates that the design of the source distribution constitutes an independent, critical axis for generative policy performance. The authors introduce LeaP, a Learnable source Prior, which replaces the standard Gaussian with a proprioception-conditioned diagonal Gaussian, designed to generate action chunks that are both observation-informed and stochastically initialized. The prior is parameterized by a lightweight MLP, jointly outputs a state-adaptive mean and variance, and is trained end-to-end with the downstream generator using a suite of objectives. Empirical evaluation across 15 RoboTwin manipulation tasks and real-robot deployments provides evidence for LeaP’s superiority, and its generality across both flow-matching and diffusion-bridge policy families.

Motivation and Prior Art

Most existing generative robot policy architectures—encompassing established works in flow-matching and diffusion modeling—fix the source for action generation as $N(0, I)$, independent of robot state or environmental context. This design forces the generator to expend capacity transporting samples from an uninformed initial state to executable, state-relevant actions. Attempts at replacing the standard source have favored deterministic, observation-conditioned initializations—derived from proprioceptive sequences (A2A), visual latents (VITA), or fused sources (BridgePolicy)—but these typically ignore uncertainty, either producing point estimates or sources with fixed, task-agnostic variance.

Importantly, recent theoretical analysis suggests that the effectiveness of generative policies arises not just from distributional expressivity but from iterative computation and stochasticity injection, yet the role of learnable, observation-adaptive source stochasticity remains underexplored. This work directly addresses whether the source should be deterministic or stochastic, and, crucially, how its parameters should be tied to observation modalities.

Methodology

LeaP Source Prior Design

LeaP replaces the observation-independent prior $q_0$ with a proprioception-conditioned diagonal Gaussian over the action chunk. The prior network—a two-layer MLP—maps the proprioceptive state feature $z_s$ to the mean $\mu_g$ and log-variance $\log \sigma_g^2$ for each action dimension. At each control step, $H$ action chunk steps are sampled via reparameterization:

$$z_{0,h} = \mu_g(z_s) + \sigma_g(z_s)\odot \epsilon_h, \quad \epsilon_h \sim N(0, I_{d_a}), \quad h = 1,\ldots,H$$

The sampled source $z_0$ initializes the generator, which performs precise, observation-conditioned refinement. Crucially, the generator and solver are unmodified: LeaP is plug-and-play, altering only the source distribution.

Training Objective

The prior and generator are trained jointly with a loss composed of:

• Negative log-likelihood (NLL): Matches the marginal of the prior to empirical action distributions.
• Symmetric Contrastive Alignment (CLIP-style): Enforces paired correspondence between sampled sources and demonstrated actions at the instance level in a learned embedding space.
• Flow Matching Loss: The standard flow-matching regression objective, which propagates gradients to both the prior and generator.

This loss structure ensures both distribution-level alignment and fine-grained sample-level grounding.

Empirical Evaluation

Simulation and Real-world Results

LeaP is assessed on 15 diverse bimanual manipulation tasks from RoboTwin 2.0 and three real-robot tasks using the Franka Research 3 platform. All methods—including strong baselines like A2A, VITA, NoPrior (vanilla baseline), and BridgePolicy—are matched in terms of observation encoders and generator architectures where possible.

LeaP achieves an average success rate of 81.6% across the RoboTwin suite, outperforming the best alternative by 6.5 to 25.5 percentage points. The most direct comparison, NoPrior (identical architecture, only source changed), lags by 25.5 points, isolating the contribution of the learnable, stochastic source prior. In real-robot deployment, LeaP remains superior (80.0% mean success), with the source prior robustness and informativeness translating to actual system performance.

Moreover, LeaP’s architectural overhead is minimal—adding only $\sim0.21$M parameters relative to $N(0, I)$0M in NoPrior—far less than additional modules in A2A or large backbone disparities in Transformer/UNet-based methods.

Efficiency and Generality

LeaP accelerates convergence and stabilizes training compared to all tested baselines. It is effective as a plug-in prior for both flow-matching and diffusion-bridge generators, improving average task success in both cases (up to +37.6 percentage points for flow matching, +8.0 for diffusion bridge). This demonstrates that observation-adaptive, learnable source design complements the choice of generative dynamics rather than being a generator-specific trick.

Ablation Studies

A comprehensive ablation isolates the pivotal design choices:

• Observation modality: Conditioning the prior solely on proprioceptive state yields the best performance. Incorporating visual features, regardless of fusion method, uniformly degrades success rates, and the prior samples are less well aligned to target actions.
• Prior parameterization: A jointly learned mean and state-adaptive diagonal variance provides the optimal trade-off; neither only-learning the mean, fixed variance, nor higher-capacity structures (full covariance, GMMs) surpass the diagonal Gaussian. Gains come not from expressivity or injected stochasticity per se, but from observation-adaptive uncertainty anchoring the source in plausible action neighborhoods.
• Supervision: Both NLL and CLIP-style alignment losses provide substantial, complementary benefit. Removing either degrades performance (from 85.3% to ~74%), and removing both reduces to 59.7%.
• Direct execution: Simply deploying the prior mean as the action (without downstream refinement) yields 0% success, confirming the necessity for generator refinement atop a well-designed prior.

Theoretical and Practical Implications

The results substantiate that source distribution design, specifically in terms of observation-adaptive stochasticity, is an underutilized but critical lever for generative policy quality. This axis is orthogonal and complementary to generator dynamics (flow, diffusion, bridges), and can be modularly improved without generator redesign. The findings suggest that priors conditioned solely on proprioceptive state maximally leverage task-context without overfitting to visual idiosyncrasies, and that learnable state-adaptive uncertainty can calibrate the generator’s initial search region without incurring parameter bloat.

Future directions include more expressive but computationally efficient structured covariances (e.g., low-rank, normalizing-flows), extension to high-DOF dexterous hands, and integration with vision-language-action policy frameworks.

Conclusion

This work establishes that the initial source distribution for generative policies is a reusable and independent design axis with significant impact on final policy efficacy. LeaP, a proprioception-conditioned, learnable diagonal Gaussian prior, is demonstrated as an effective and general plug-in that consistently outperforms strong deterministic and stochastic baselines across a wide range of manipulation tasks and model classes. The modularity, negligible parameter overhead, and rapid convergence of LeaP suggest that source prior design should be considered a first-class citizen in the architecture of generative robot policies, especially for scalable real-world deployment.

Reference: "Where Should Action Generation Begin? A Learnable Source Prior for Generative Robot Policies" (2606.17408)