Diffusion Policy Policy Optimization

In plain words

A robot policy can learn from human demonstrations. The usual recipe is to copy the demonstrations, which often leaves the robot a little worse than the demonstrator. The paper asks whether the copied policy can be improved by trial and error, the way reinforcement learning improves policies.

The policy in question is a diffusion model. It starts from random noise and removes it in several small steps until it reaches an action. Earlier work assumed that trial-and-error updates do not work well on this kind of policy, because the probability of an action is hard to compute. The paper shows that each small noise-removal step does have a tractable probability, so the usual policy-gradient update can be applied to all of them. The authors call the method DPPO. It fine-tunes a policy that was pre-trained from demonstrations, and it works in simulation and on a real robot arm.

TL;DR

• DPPO fine-tunes a behaviour-cloned diffusion policy with PPO §4.2.
• The denoising chain is written as an inner MDP nested in the environment MDP, so each denoising step has a Gaussian likelihood §4.1 eq. 4 eq. 6.
• Practical choices matter: fine-tune only the last few denoising steps, or use DDIM; clip the noise schedule from below for exploration §4.3.
• Author-reported: strongest fine-tuning results among the compared RL methods on the benchmarks used; first RL method reported to reach >50% on Robomimic Transport from state or pixels §5.3.
• The official code is MIT-licensed and matches the paper in most hyperparameters, with a few differences listed below.

Problem

Behaviour cloning from a small or noisy demonstration set gives policies that are limited by the data §1. Fine-tuning them with reinforcement learning is attractive, but for diffusion policies the action likelihood needed by policy gradients is not available in closed form §1 §3. Earlier fine-tuning methods for diffusion policies are off-policy Q-learning, weighted regression, or indirect methods §2 §5.1. The paper's question is whether on-policy policy-gradient fine-tuning can work for a diffusion policy §1.

Method

Setting. Policy \pi_\theta is a DDPM-style diffusion policy that outputs a chunk of T_p future actions, of which T_a \le T_p are executed §3. The denoising chain has K steps §3 eq. 3.

Two-layer MDP (Diffusion Policy MDP). Each environment step contains an inner MDP over the K denoising steps §4.1 fig. 3. Inner states carry the environment state and the partially denoised action; inner actions are the denoised action after one step. Reward is only given at the last denoising step (k=0) §4.1 eq. 4 eq. 5. The denoising step is a Gaussian, so its likelihood is analytic eq. 6. Policy gradient in this outer-plus-inner MDP is eq. 7.

PPO instantiation. DPPO applies the clipped PPO objective eq. 8 to the two-layer MDP §4.2. The advantage uses an environment-discounted return with discount \gamma_{\text{env}} eq. 9, then is multiplied by a denoising discount \gamma_{\text{denoise}}^k eq. 10. The value function depends only on the environment state s_t, not on the partly denoised action eq. 11 §5.5. The paper uses GAE for the advantage estimate in practice [§4.2 footnote].

Best practices §4.3:

1. Fine-tune only the last K' denoising steps. A frozen copy of the pre-trained weights handles the earlier steps §4.3.
2. Fine-tune DDIM sampling instead, with K^{\text{DDIM}} \ll K steps, e.g. 5. Training uses \eta=1 (DDPM-equivalent) and evaluation uses \eta=0 eq. 12 §4.3.
3. Clip the noise schedule from below at \sigma^{\text{exp}}_{\min} when sampling, for exploration. Clip at \sigma^{\text{prob}}_{\min}=0.1 when evaluating the Gaussian likelihood §4.3.
4. Architecture: MLP is more stable to fine-tune. UNet allows fine-tuning with T_a < T_p §4.3.

Algorithm 1 gives the full loop alg. 1 app. B.

# DPPO: fine-tune a behaviour-cloned diffusion policy  [alg. 1] [app. B]
# Input:  pi_theta pre-trained with the BC loss  [eq. A1]
# Output: fine-tuned pi_theta

def dppo_finetune(policy, value, envs, num_iterations):
    # Each iteration starts from a frozen copy of the current policy [alg. 1]
    for iteration in range(num_iterations):
        old_policy = copy(policy)
        buffer = []
        # Roll out N environments in parallel for T steps each  [alg. 1]
        for env_step in range(num_env_steps):
            # Denoise K times, from pure noise to the action chunk  [§4.1]
            for k in range(K - 1, -1, -1):
                # Sample one denoising step from the frozen policy  [eq. 6]
                action_k = sample_denoising_step(old_policy, obs, k)
                # Reward is only received at the last denoising step  [eq. 4]
                reward = env_reward if k == 0 else 0.0
                buffer.add(k, obs, action_k, reward)
            # Execute the final denoised action in the environment  [alg. 1]
            obs = step_env(envs, action_k)
        # Advantage at environment steps, via GAE  [eq. A2]
        advantages = gae(buffer, value, lam=GAE_LAMBDA, gamma=GAMMA_ENV)
        # Several minibatch updates per iteration  [alg. 1] [app. C.10]
        for update in range(num_updates):
            batch = sample_minibatch(buffer)
            # Denoising discount: earlier (noisier) steps get smaller weight  [§4.2]
            discounted = GAMMA_DENOISE ** batch.k * batch.advantages
            # Clipped PPO objective on each denoising step  [eq. A3]
            policy_loss = ppo_clip_loss(policy, old_policy, batch, discounted)
            # Value regression on environment-state returns  [eq. A4]
            value_loss = squared_error(value(batch.env_states), batch.returns)
            update_parameters(policy, value, policy_loss, value_loss)
    return policy

Note that the pseudocode above is a sketch. The exact order of rollout and buffer bookkeeping is in alg. 1.

Hyperparameters

Values are from the paper's appendix app. C.10 unless marked. Where the appendix gives two values for the same task (see Pitfalls), both are listed.

Diffusion-based RL comparison (state input), DPPO row

Common to all: critic learning rate 1e-3, critic MLP [256, 256, 256], and pre-training with the Gym/Robomimic recipe below app. C.10.

Pre-training (Gym): 3000 epochs, batch 128, learning rate 1e-3 decayed to 1e-4 with cosine schedule, weight decay 1e-6, EMA decay 0.995. Observations and actions normalised to [0,1] by dataset min/max app. C.2.

Pre-training (Robomimic): batch 128, learning rate 1e-4 decayed to 1e-5 with cosine schedule, weight decay 1e-6. Diffusion policies 8000 epochs, Gaussian and GMM 5000 epochs app. C.6. Robomimic Multi-Human data is the default; Proficient-Human is used in some experiments §5.1 app. C.6.

Fine-tuning (Gym): 40 parallel MuJoCo environments; evaluation with \sigma^{\text{exp}}_{\min}=0.001 for 40 episodes app. C.2.

Fine-tuning (Robomimic): 50 parallel environments; evaluation with \sigma^{\text{exp}}_{\min}=0.001 for 50 episodes app. C.6.

Policy-parameterisation comparison, Robomimic state input app. C.10:

Policy-parameterisation comparison, Robomimic pixel input app. C.10: T_a 4, 4, 8 for Lift/Can, Square, Transport; DPPO-ViT-MLP uses K=100, K'=5 (DDIM), \gamma_{\text{denoise}}=0.9, \sigma^{\text{exp}}_{\min} 0.1, 0.1, 0.08, \sigma^{\text{prob}}_{\min}=0.1, \epsilon=0.01 (annealed), batch 37500, 50000, 50000. Gaussian-ViT-MLP uses \sigma_{\text{Gau}} 0.1, 0.1, 0.08 and batch 7500, 10000, 10000 app. C.10.

FurnitureBench (sim, and sim-to-real) app. C.8 app. C.10:

Hardware (sim-to-real, One-Leg) app. C.8: Franka Panda, differential IK with dt = 1/10 s, joint impedance control at 1 kHz via Polymetis. Four RealSense D435 cameras, AprilTag pose estimation, exponential filter \alpha = 0.25, rejection if a jump exceeds 5 cm or 20°. Domain randomisation adds zero-mean noise with 0.03 std to actions app. C.8.

Diffusion RL baselines and Gaussian/GMM/Transformer baselines are in app. C.3, app. C.4, app. C.7 and app. C.10.

Results

Author-reported. Setup: each task as described in Hyperparameters; pre-trained from demonstrations, then fine-tuned with RL; success rate and reward curves averaged over 3 or 5 seeds (figures; exact numbers are not in the text for Gym).

Wall-clock times: on Robomimic state input, DPPO-MLP is 21–24% slower per iteration than the Gaussian and GMM baselines, and DPPO-UNet is about 49% slower than DPPO-MLP app. E. Per-iteration times for Gym are in app. E (DPPO is about 41% faster than DAWR, 37% faster than DIPO and 12% faster than DQL on average, per the text).

Paper vs code

Repository: irom-lab/dppo at commit cc7234ad7ff39a8f32de3af903606723a16f0648, MIT licence. The paper's abstract links diffusion-ppo.github.io, which links this repository. The repository belongs to the Intelligent Robot Motion (IRoM) lab, whose members include author Anirudha Majumdar.

Checked:

• Advantage normalisation within each batch is done in the code: after the advantage estimate, advantages = (advantages - mean) / (std + 1e-8) diffusion_ppo.py. The paper describes advantages only as GAE with the denoising discount §4.2 app. B.
• The denoising discount \gamma_{\text{denoise}}^{K-1-i} is applied to the advantage in the same way as eq. 10 diffusion_ppo.py.
• The PPO ratio uses the clipped form of eq. A3. The paper's clip schedule (ratio 0.1\epsilon at the noisiest step, exponential in between) is implemented as an exponential interpolation between a base and the main \epsilon diffusion_ppo.py.
• The value loss has a factor of 0.5 and an optional value clip. The paper's eq. A4 has no factor of 0.5 diffusion_ppo.py. This is a constant factor, so it changes the effective value learning rate but not the minimiser.
• Batch size 7500 and KL target 1 for Lift are as in the paper's diffusion-comparison table ft_ppo_diffusion_mlp.yaml. Pre-training uses 20 denoising steps pre_diffusion_mlp.yaml.
• \sigma^{\text{prob}}_{\min}=0.1 and \sigma^{\text{exp}}_{\min}=0.1 (Hopper config) ft_ppo_diffusion_mlp.yaml.

Differences or things the paper does not mention:

• The code clamps log-probabilities to [-5, 2] before computing the ratio diffusion_ppo.py. The paper does not mention this clamp.
• The loss only uses the first reward_horizon action entries (default act_steps) of the chunk when computing the policy-gradient term. The code comment says this is done to backpropagate only through the steps executed in the environment diffusion_ppo.py train_ppo_diffusion_agent.py. The paper describes the full chain of K' denoising steps and does not state this restriction.
• The total loss adds an entropy term weighted by ent_coef, a BC term weighted by bc_loss_coeff, and an optional learnable noise scale eta train_ppo_diffusion_agent.py train_ppo_diffusion_agent.py. The paper does not describe these. Default values for ent_coef, bc_loss_coeff and learn_eta were not checked for every task in this compilation.
• The Hopper config sets clip_ploss_coef_base equal to clip_ploss_coef (both 0.01), which makes the denoising-dependent clip schedule flat ft_ppo_diffusion_mlp.yaml. The paper's ε schedule (0.1 ε at the noisiest step) is therefore not active in that config. Other task configs were not checked.
• The code's GAE loop runs over environment steps only, as in the paper train_ppo_diffusion_agent.py.
• The code reports the approximate KL with the k3 estimator and stops the batch loop when it exceeds target_kl diffusion_ppo.py train_ppo_diffusion_agent.py. The paper says this threshold is 1 and is never hit in practice app. B.
• The paper's sim-to-real domain randomisation (0.03 std action noise, state noise matched to hardware) app. C.8 is not checked here against the FurnitureBench configs.

Pitfalls

• Tables and figures in the appendix conflict on batch size. The diffusion-comparison table uses batch 7500 for Lift/Can, 10000 for Square and Transport app. C.10. The policy-comparison table uses 75000, 100000, 100000 for the same tasks app. C.10. The code's Lift config uses 7500 ft_ppo_diffusion_mlp.yaml. Start from the code value and check the paper's intent.
• Gym results are only in figures. The text gives no numeric success or return values for Gym §5.1 fig. 5. Do not invent targets from the curves.
• Batch size scales with denoising steps. The paper uses a batch of roughly the Gaussian batch size times the number of fine-tuned denoising steps app. B. Each rollout step produces K entries, so memory grows with K' app. D.
• Exploration noise is sensitive. Square collapses without \sigma^{\text{exp}}_{\min}=0.1 app. D. FurnitureBench needs \sigma^{\text{exp}}_{\min}=0.04 for the pre-trained policy to have nonzero success app. C.8.
• Likelihood clipping. Clipping \sigma_k at 0.1 when computing likelihood is needed for stability §4.3. The paper does not say how sensitive results are to this value beyond the sweep in app. D.
• Action chunk size changes which policies train. UNet can fine-tune with T_a < T_p; MLP with T_a < T_p is unstable §4.3 [footnote in app. B].
• Pixel results use DDIM. Pixel and furniture runs fine-tune 5 DDIM steps with \eta=1 during training, then \eta=0 at evaluation §4.3 app. C.10.
• Cost. Each environment step costs K network evaluations. Training from scratch with 10 steps was about 6× slower than Gaussian app. D. Fine-tuning On furniture, DPPO-UNet is about 20% slower than Gaussian-MLP app. E.
• Sim-to-real claims rest on one task and 20 trials §5.4 app. C.8.
• Equation and figure numbers in this card were counted from the LaTeX source, not from the rendered PDF. If a citation looks off by one, check the PDF.

Lineage

• Builds on Diffusion Policy (Chi et al., 2023), the behaviour-cloned diffusion policy DPPO fine-tunes §3.
• Uses PPO (Schulman et al., 2017) and GAE for the RL update §4.2.
• Compared against: Q-learning and weighted-regression diffusion RL methods (DIPO, IDQL, DQL, QSM, DRWR, DAWR) §5.1; demo-augmented RL methods (RLPD, Cal-QL, IBRL) §5.2; Gaussian, GMM and transformer policies §5.3.
• Related: policy gradient on the denoising process of text-to-image diffusion (Black et al., 2023; Clark et al., 2023) §1 §4.1.