Sample-adaptive Lipschitz budget for diffusion policies

(1) Integration point: replace the unconstrained drift MLP inside a conditional diffusion policy with a spectrally normalized drift network. The network receives diffusion time t, observation x, and current action state A_t, and predicts b_theta(t,x,A_t); this modification affects both the training-time diffusion loss and inference-time Euler or Heun sampling. Use a fixed candidate budget set K_grid, for example {0.5, 1, 2, 4, 8, 16}, after normalizing state and action coordinates.

(2) Pseudocode:

for K in K_grid:
    construct drift b_theta with Wbar_l = W_l * min(1, c_l / sigma_max(W_l))
    choose c_l so product_l c_l = K / product_l Lip(activation_l)
    train b_theta on training trajectories
    for validation state x:
        A = sample_base_noise()
        for j in range(N_steps):
            t = j*T/N_steps
            A = A + b_theta(t, x, A)*dt + sigma*sqrt(dt)*normal()
        evaluate validation return or action log-likelihood
select K with the best validation score

The displayed Lipschitz inequality is the target constraint and the product of spectral norms is the computable upper bound. Estimate each sigma_max(W_l) with one to three power iterations per optimizer step; estimate actual local sensitivity separately with random finite differences, K_local = max ||b(z+epsilon*u)-b(z)||/epsilon over validation samples. The approximation term and sample-size rate come from the paper; the practical K choice is selected by validation because constants, action dimension, horizon, and noise level are unknown. First experiment: train a small two-layer diffusion policy on a multimodal two-dimensional continuous-control imitation or offline-RL task, comparing unconstrained training, fixed spectral normalization, and the K grid at equal parameter count and sampling steps. Measure terminal-action Wasserstein distance to demonstrations, return, validation likelihood, empirical K_local, and gradient stability. Success is a U-shaped validation-error or return curve in K, with small K underfitting multimodality and large K overfitting, plus improved held-out return at the selected K over an unconstrained baseline.