Condition-Dependent Diffusion Policies

• Condition-dependent priors in diffusion policies are advanced techniques that adapt the initial noise distribution based on task context, enhancing model specificity.
• They leverage learned Gaussian and non-isotropic priors to capture task-relevant features, thereby guiding the diffusion process with contextual biases.
• Empirical evaluations show these methods improve sample efficiency and robustness in applications such as robot learning, planning, and generative modeling.

Condition-dependent priors in diffusion policies refer to the class of techniques in which the probabilistic starting point (or noise distribution) of the diffusion process is adapted according to the task, context, or observation. This contrasts with conventional approaches using fixed, unconditional priors—typically standard normal distributions—at the highest-noise step. Recent advances have demonstrated that such condition dependence in priors substantially impacts both the expressivity and efficiency of diffusion-based policies, enabling improved robustness, sample efficiency, and task specificity across robot learning, planning, and generative modeling domains.

1. Formal Definition and Theoretical Foundations

In diffusion policies, the generative process iteratively denoises samples initialized from a prior distribution defined at the maximal noise timestep. For standard approaches, this prior is unconditional,

$p(x_T) = \mathcal{N}(x_T; 0, I),$

yielding joint or conditional generative models by subsequent conditioning in the denoising network. Condition-dependent priors replace $p(x_T)$ with a distribution $p_\phi(x_T \mid s)$, where $s$ is a task, observation, or environment context. Typical parameterizations include Gaussian distributions with context-dependent means and variances or, more generally, mixture models or GP-based structures.

The principal motivation is to encode informative inductive biases, focus the initial sample mass in regions relevant for the task, or disambiguate highly multimodal or weakly identifiable solutions. Theoretical results show that such context anchoring improves gradient informativeness in conditional training and can alleviate loss collapse due to insufficient condition separation, as formalized in, e.g., Cocos (Dong et al., 16 May 2025).

2. Representative Architectures and Mathematical Formulations

Distinct condition-dependent diffusion frameworks include:

• Learned Gaussian Priors: In prior-guided planning, $$p_\phi(x_T|s) = \mathcal{N}(x_T; \mu_\phi(s), \text{diag}(\sigma_\phi^2(s)))$$, where $\phi$ is an encoder (e.g., GRU, transformer) of the environment state. The prior is trained to maximize expected trajectory return regularized by KL divergence to an unconditioned baseline, e.g., (2505.10881).
• Non-isotropic Task-Aware Priors: Task-structured priors leverage means and covariances from Gaussian Process Motion Planning (GPMP), conditioning the mean and covariance on task-defining key states, as in $$q(\tau^i | \tau^{i-1}) = \mathcal{N}(\sqrt{\alpha_i}\tau^{i-1} + (1-\sqrt{\alpha_i})\mu, (1-\alpha_i)K)$$, with $(\mu, K)$ task- or context-conditioned (Kim et al., 30 Sep 2025).
• Explicit Plug-and-Play Priors: A fixed pretrained diffusion prior $p(x)$ is combined with an external constraint $c(x, y)$ to define $p(x|y) \propto p(x)\cdot c(x, y)$, where $c(x, y)$ can encode arbitrary differentiable relationships (Graikos et al., 2022).
• Product-of-Experts Compositional Scores: Multiple diffusion models trained on different tasks or context segments are combined in score space to form a newer, composed, condition-dependent prior by $$q^{\mathrm{comp}}(a_0) \propto \prod_{i} q_i(a_0)^{w_i}$$, with $w_i$ learned (Patil et al., 2024).

These models require different algorithmic and statistical handling for the forward noising process, the reverse kernel, inference, and optimization objectives.

3. Training Procedures and Loss Functions

Commonly, the training objective adapts the denoising loss to respect the condition-dependent prior. For learned Gaussian priors,

$$\mathcal{L}_{\text{prior}}(\phi) = -\mathbb{E}_{s, x_T}\left[\overline{V}_\theta(x_T)\right] + \alpha\,\mathrm{KL}(p_\phi(\cdot|s)\,||\,\mathcal{N}(0,I))$$

is used, where $\overline{V}_\theta$ is a latent value function estimating denoised trajectory return (2505.10881). For factored policies with prioritization, the total training loss involves sequential fitting:

• Fit $\varepsilon_\mathrm{base}$ on prioritized modalities alone,
• Freeze $\varepsilon_\mathrm{base}$, then fit $\varepsilon_\mathrm{res}$ on all modalities with loss:

$$\mathcal{L}_\mathrm{res} = \mathbb{E}\left[\frac{1}{2}\| x_0 - \varepsilon_\mathrm{base}(x_t, y^{(k)}, t) - \varepsilon_\mathrm{res}(x_t, y^{(M)}, t)\|^2\right]$$

—obtaining an additive decomposition of the score (Patil et al., 20 Sep 2025).

For hierarchical noise models and non-isotropic priors, training minimizes a Mahalanobis-structured MSE loss, accounting for the covariance structure $K$ in the data (Kim et al., 30 Sep 2025).

4. Algorithmic Implementation and Inference

Sampling from a condition-dependent prior typically involves two main changes to the diffusion procedure:

• Initializing samples $x_T \sim p_\phi(x_T|s)$ with the learned or structured prior rather than $\mathcal{N}(0, I)$.
• Adapting the denoising kernels to accommodate the prior’s mean and (if needed) covariance:
  • For structured covariances, closed-form expressions ensure consistency during both training and sampling.
  • Compositional models combine multiple scores/denoisers at each step with optimized convex weights.

Pseudocode for these steps generally mirrors standard diffusion rollouts with the aforementioned modifications in the terminal time and score computation. For sequentially factored policies, the approach runs frozen $\varepsilon_\mathrm{base}$ and then adds the (adaptively trained) $\varepsilon_\mathrm{res}$ term at each iteration (Patil et al., 20 Sep 2025). Hierarchical models, as in task-conditioned motion planners, instantiate the (mean, covariance) prior hierarchy before running the diffusion chain (Kim et al., 30 Sep 2025).

5. Empirical Impact and Evaluation

Condition-dependent priors consistently yield statistically significant gains in metrics of sample efficiency, robustness, and alignment with task or data constraints:

• Sample Efficiency: FDP improves success by 10-20% absolute in low-data regimes relative to standard joint-conditioning, especially pronounced at 10-50 demonstration counts (Patil et al., 20 Sep 2025).
• Robustness: Policies with factored or structured priors maintain high success rates under severe distribution shifts (visual distractors, occlusions), with up to 40% absolute gain vs. catastrophic baseline failure (Patil et al., 20 Sep 2025).
• Compositional Flexibility: PoE-based condition-dependent priors via DSE reduce data-distribution discrepancy by 30–50% (MMD-FK) over baseline finetuning and naive composition in few-shot imitation (Patil et al., 2024).
• Reduced Loss Collapse: In large-scale multimodal settings, condition-dependent source priors prevent degradation of condition signal and dramatically accelerate convergence (94.8% average task success with condition-dependent prior vs. 86.5% without) (Dong et al., 16 May 2025).

Structured priors (e.g., GPMP) concentrate sample mass on dynamically feasible and task-relevant trajectories, as evidenced by success rates tripling in goal-reaching and stacking tasks (Kim et al., 30 Sep 2025). Behavior-regularized and bridge-based priors improve performance when function-evaluation budgets are very limited, decreasing inference cost by up to $10\times$ versus online planning (2505.10881, Srivastava, 2024).

6. Modality Prioritization and Factorization

Factorized Diffusion Policies (FDP) introduce explicit prioritization over input observation modalities (proprioception, vision, tactile, etc.) by decomposing the score function using Bayes' rule:

$$\nabla_{x_t} \log p(x_t|y^{(M)}) = \nabla_{x_t} \log p(x_t|y^{(k)}) + \nabla_{x_t} \log p(y^{(k+1:M)}|x_t, y^{(k)})$$

A two-stage parameterization trains a base denoiser $\varepsilon_\mathrm{base}$ on prioritized modalities, then a residual denoiser $\varepsilon_\mathrm{res}$ to correct for de-prioritized modalities. Empirically, this design yields greater sample efficiency and robustness, and allows easy adaptation to different observability or sensor-reliability regimes, provided the prioritization is chosen to match task-structure (Patil et al., 20 Sep 2025).

7. Limitations and Open Directions

Despite their demonstrated benefits, condition-dependent priors are subject to several structural and practical limitations:

• Choice and Optimization of Conditioning: The search space over modalities or composition weights can grow combinatorially with the number of potential inputs, though ablations suggest only a few orderings are empirically relevant (Patil et al., 20 Sep 2025).
• Static Prioritization: Current factored approaches fix the ordering of prioritized modalities for the entire diffusion rollout; endogenously variable or state-adaptive prioritization remains an active research question (Patil et al., 20 Sep 2025).
• Computational Overhead: Marginal increases in inference latency (e.g., 1.5$\times$ for factored policies) are incurred, but remain tractable ($\sim$150 ms/action) (Patil et al., 20 Sep 2025).
• Scalability in Planning: Approximations such as those required for Schrödinger bridge solvers can reduce expressivity relative to standard DDPMs with sufficient steps, potentially constraining applicability in high-dimensional or online control scenarios (Srivastava, 2024).

A promising implication is that improved priors—especially ones encoding task structure or learned from related policies—can reduce the amortized cost of training and inference, offering a principled strategy for scalable, robust, and adaptive diffusion policies for control, planning, and generative tasks.