Technical Summary: Parameterized Diffusion Policies (PDP)

Problem Statement

Real-world robotic tasks are often underspecified, where a single goal can be achieved by multiple distinct action sequences (multimodal behaviors). While standard Diffusion Policies (DP) have emerged as a powerful paradigm for modeling these multimodal distributions, they suffer from a lack of controllability. Specifically, standard evaluations often focus on observation-side distribution shifts (e.g., changes in camera appearance), but deployment failures frequently arise from constraint-induced behavior shifts. In these scenarios, environmental changes (e.g., new obstacles or contact constraints) invalidate previously demonstrated strategies, forcing the robot to select a different existing mode or discover a novel behavior.

Current diffusion policies rely on stochastic noise for diversity. Steering these policies at test time by optimizing the high-dimensional noise space is often ill-conditioned: small perturbations in noise can lead to large, unpredictable, or infeasible trajectory changes. Furthermore, existing methods for behavior steering either require retraining policy weights (slow, data-heavy) or rely on discrete skill tokens that restrict control to a finite set of behaviors, preventing smooth interpolation.

Methodology: Parameterized Diffusion Policy (PDP)

The authors propose Parameterized Diffusion Policy (PDP), a framework that transforms diffusion from a mechanism for stochastic diversity into a precise, optimizable tool for behavior steering. PDP conditions the diffusion denoiser on a low-dimensional, continuous behavior latent variable $z$, learned from demonstrations.

1. Geometry-Aligned Behavior Manifold

The core innovation is the construction of a behavior manifold where Euclidean distances between latent codes reflect the semantic similarity of physical trajectories.

• Trajectory Encoder: A trajectory encoder $E_\phi$ maps demonstration trajectories $\tau$ to a latent code $z$.
• Geometry Alignment: Unlike standard Variational Autoencoders (VAEs), PDP explicitly aligns the latent space geometry with physical trajectory similarity using Soft-DTW (Soft Dynamic Time Warping). The training objective includes a geometry loss $L_{geo}$ that enforces a linear correspondence between the latent distance $\|z_i - z_j\|$ and the physical trajectory distance computed via Soft-DTW. This ensures that small changes in $z$ induce smooth, predictable changes in the executed trajectory.
• Joint Objective: The encoder-decoder is trained to minimize reconstruction loss, KL-divergence (for regularization), and the geometry alignment loss.

2. Parameterized Diffusion Denoiser

The diffusion policy $\pi_\theta$ generates action chunks conditioned on both the environmental context $c_t$ and the behavior latent $z$.

• Global Modulation: Instead of simple concatenation, the latent $z$ is integrated into the denoiser via a global modulation mechanism (inspired by FiLM). The latent is transformed into layer-specific affine parameters ($\gamma, \beta$) that scale and shift intermediate feature maps. This deep integration allows the low-dimensional latent to effectively steer the high-dimensional action generation process.
• Decoupled Training: The encoder is trained with a stop-gradient operator during the denoiser optimization phase to prevent the diffusion loss from distorting the carefully structured geometry of the latent manifold.

3. Test-Time Control and Adaptation

PDP enables fast, data-light adaptation without updating policy weights.

• Latent Fitting: When faced with new constraints or a target behavior $\tilde{\tau}$, the system optimizes the latent code $z$ to minimize a diffusion-consistent fitting loss. This identifies the specific behavior mode that best satisfies the new constraints.
• Semantic Warm-Start: The optimization is initialized using the encoder's prediction on the new demonstration ($z_0 = E_\phi(\tilde{\tau})$), placing the search in a well-conditioned region of the manifold.
• Long-Horizon Adaptation: For complex tasks, PDP supports piecewise-constant latent sequences, allowing different phases of a task to be adapted independently.

Key Contributions

1. Parameterized Diffusion Framework: Introduces PDP, which replaces stochastic noise steering with a continuous, geometry-aligned latent parameter $z$ for explicit behavior control.
2. Geometry-Aligned Manifold: Demonstrates that aligning latent distances with physical trajectory similarity (via Soft-DTW) enables smooth interpolation and predictable steering, solving the ill-conditioned nature of noise-space optimization.
3. Deep Latent Integration: Proposes a global modulation architecture that successfully integrates low-dimensional behavior signals into high-dimensional diffusion denoisers, outperforming naive concatenation.
4. Efficient Test-Time Adaptation: Shows that adapting only the latent $z$ (via gradient descent) allows for rapid adaptation to novel constraints using a single demonstration, without retraining the policy.

Experimental Results

The authors evaluated PDP on a suite of nine multimodal manipulation domains in both simulation (RLBench, Robomimic, Franka Kitchen, D3IL) and on a real Franka Emika Panda robot.

• Original Multimodal Performance: PDP achieved consistently high success rates across all tasks, matching or exceeding state-of-the-art baselines (including standard DP, BC-GMM, and VQ-BeT). It demonstrated superior stability in reconstructing multimodal distributions compared to standard DP, which often produced noisy, inconsistent rollouts due to mode ambiguity.
• Constraint-Induced Shifts: In scenarios where environmental constraints invalidated most training modes:
  • Existing Mode Fitting: PDP achieved near-perfect success (e.g., 100% on CloseDrawer), significantly outperforming baselines that failed to converge to the single feasible mode.
  • Novel Behavior Discovery: In "zero-mode-feasible" settings (where no training mode works and a new behavior must be synthesized), PDP was the only method to consistently succeed. Baselines, including those that fine-tuned weights (BC) or used evolutionary search (Diff-ES), failed to synthesize feasible strategies absent from the training data.
• Real-Robot Robustness: On the real robot OpenDrawer task, PDP maintained a 100% success rate under constraint shifts, whereas standard DP's success dropped to 20%. This highlights PDP's robustness to real-world noise and its ability to provide a stable control handle.
• Ablation Studies:
  • Removing geometry alignment or using discrete latents caused severe performance degradation, confirming the necessity of a continuous, metric-aligned manifold.
  • Global modulation significantly outperformed concatenation, reducing training loss and improving adaptation success.
  • Semantic warm-starts were critical; random initialization failed to find viable strategies even after extensive optimization steps.

Significance and Claims

The paper claims that PDP fundamentally shifts the paradigm of diffusion-based control from "noise-driven diversity" to "parameterized behavior steering." By structuring the latent space to reflect physical trajectory similarity, PDP transforms diffusion into a tool that is not only expressive but also controllable and optimizable.

The authors emphasize that this approach addresses a critical gap in robot learning: the ability to adapt to changing constraints rapidly and reliably without retraining. They argue that while diffusion models are excellent at representing multimodal data, their utility in real-world deployment is limited without a structured interface for behavior selection. PDP provides this interface, enabling robots to:

1. Smoothly interpolate between known strategies to synthesize novel, feasible behaviors.
2. Rapidly adapt to new constraints by optimizing a compact set of parameters (the latent $z$) rather than millions of network weights.
3. Operate robustly in real-world settings where execution noise and constraint shifts are common.

The work suggests that explicit, geometry-aligned behavior manifolds are essential for the next generation of adaptable robotic policies, particularly in scenarios requiring the synthesis of behaviors not explicitly present in the training dataset.