PaperlayerView exact revision on arXiv

Source-linked AI summary

Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation

Aaditya Prasad, Kevin Lin, Jimmy Wu, Linqi Zhou, Jeannette Bohg

arXiv:2405.07503v2cs.ROcs.AI

TL;DR

Diffusion Policy provides strong robot-control performance but is slow and computationally demanding for dynamic or resource-constrained systems. This paper distills a pretrained diffusion policy into Consistency Policy using self-consistency along learned trajectories, achieving roughly order-of-magnitude faster inference with competitive success rates. The method’s main trade-off is reduced action-distribution multimodality relative to Diffusion Policy.

  • Problem

    Diffusion Policy’s sequential denoising requires costly inference, limiting its use in dynamic tasks and on robots with constrained onboard computation.

  • Method

    Consistency Policy distills a pretrained Diffusion Policy by enforcing consistency along learned trajectories for single- or few-step generation.

  • Results

    Across 9 simulation and real-world tasks, Consistency Policy increases inference speed dramatically without sacrificing success rates.

  • Takeaways & Limitations

    Robustness to teacher quality can reduce the need for extensive teacher-model testing, while selected consistency and sampling design choices support performance.

  • Takeaways & Limitations

    Consistency Policy trades off some of Diffusion Policy’s multimodality, with both the teacher EDM policy and student favoring one side of Push-T.

Abstract

from arXiv · show

Many robotic systems, such as mobile manipulators or quadrotors, cannot be equipped with high-end GPUs due to space, weight, and power constraints. These constraints prevent these systems from leveraging recent developments in visuomotor policy architectures that require high-end GPUs to achieve fast policy inference. In this paper, we propose Consistency Policy, a faster and similarly powerful alternative to Diffusion Policy for learning visuomotor robot control. By virtue of its fast inference speed, Consistency Policy can enable low latency decision making in resource-constrained robotic setups. A Consistency Policy is distilled from a pretrained Diffusion Policy by enforcing self-consistency along the Diffusion Policy's learned trajectories. We compare Consistency Policy with Diffusion Policy and other related speed-up methods across 6 simulation tasks as well as three real-world tasks where we demonstrate inference on a laptop GPU. For all these tasks, Consistency Policy speeds up inference by an order of magnitude compared to the fastest alternative method and maintains competitive success rates. We also show that the Conistency Policy training procedure is robust to the pretrained Diffusion Policy's quality, a useful result that helps practioners avoid extensive testing of the pretrained model. Key design decisions that enabled this performance are the choice of consistency objective, reduced initial sample variance, and the choice of preset chaining steps.

I. INTRODUCTION

Diffusion Policy offers strong imitation-learning performance but requires costly sequential denoising, limiting deployment in dynamic or compute-constrained settings. Consistency Policy addresses this through distillation aimed at much faster inference while maintaining competitive performance.

  • Diffusion Policy achieves state-of-the-art imitation-learning performance across varied robotics tasks.
  • 100 denoising steps can take around one second per action generation on an NVIDIA T4.Reducing evaluations can degrade performance.
  • Slow inference constrains Diffusion Policy in dynamic tasks and on robots with limited onboard computation.Quasi-static tasks may tolerate lengthy reaction times, whereas balancing and dynamic navigation often require faster control frequencies.
  • Consistency Policy distills a pretrained diffusion model by enforcing consistency between denoising points along learned ODE trajectories.Consistency-based distillation supports larger denoising steps and can enable single- or few-step generation.
  • Consistency Policy is about an order of magnitude faster than the fastest baseline while maintaining similar or higher success rates across varied tasks.The evaluation compares it with Diffusion Policy and inference-acceleration methods; parallel sampling can increase memory requirements and remain slower than single-step prediction.

III. CONSISTENCY POLICY

The paper formulates Consistency Policy as a Consistency Trajectory Model built on diffusion-model denoising trajectories. It uses the diffusion process to map noisy actions toward observation-conditioned expert actions and supports single- or few-step inference.

  • Consistency Policy is formulated as a Consistency Trajectory Model and distilled from a trained teacher Diffusion Policy.Inference can use either a single step for maximum speed or three steps for greater accuracy.
  • A. Preliminaries: Diffusion models map Gaussian-noise actions x_T to expert actions x_0 conditioned on the current observation.The diffusion time t parameterizes the path from the simple Gaussian distribution at T to the expert action distribution at 0.
  • A. Preliminaries: The Probability Flow ODE describes forward noising and backward denoising, with a fully denoised action serving as the policy prediction.The paper uses “trajectory” for this ODE path, not for robot motions or demonstrated actions.
  • A. Preliminaries: The PFODE uses a drift coefficient, diffusion coefficient, and time-dependent noised action distribution to describe denoising dynamics.The score ∇log p_t(x_t|o) is approximated with a neural network because it is generally intractable to compute directly.
  • A. Preliminaries: Forward diffusion constructs noisy training points by perturbing expert data with Gaussian noise at selected diffusion times.The perturbed distribution is the expert action distribution convolved with N(0, t^2I).

B. Training

Training first fits an EDM teacher to estimate denoising dynamics, then distills it into a student suitable for few- or single-step generation. Numerical integration and the DSM objective define the teacher-training procedure.

  • Few- or single-step generation is trained by first fitting a teacher model and then distilling it into a student model.
  • 1) Teacher Model (EDM): The EDM teacher takes x_t, diffusion time t, and conditioning o as inputs to estimate the PFODE trajectory derivative.Numerical integration is required to convert repeated derivative estimates into positions along the trajectory.
  • 1) Teacher Model (EDM): Repeated derivative estimation and numerical integration cause Diffusion Models’ slow inference speed.
  • 1) Teacher Model (EDM): The DSM objective trains the EDM model to predict the ground-truth initial position x_0 from a sampled trajectory point (x_t, t).The paper uses the pseudohuber loss as the distance metric.
  • 1) Teacher Model (EDM): The pseudohuber metric bridges standard l_1 and l_2 norms and handles outliers more effectively than the original EDM l_2 loss.The implementation sets its small constant c to 0.00054.
  • 1) Teacher Model (EDM): Teacher-model trajectory integration uses Heun’s second-order solver with the timestep discretization scheme from EDM.

2) Student Model (Consistency Policy)

Consistency Policy distills a teacher Diffusion Policy by enforcing self-consistency between denoising paths along the same PFODE. The student predicts earlier states, and its consistency loss is combined with DSM loss.

  • The student gθ(xt, t, s; o) takes a PFODE position, diffusion time, target time, and observation, then estimates xs.
  • CTM compares student denoising paths from different PFODE points to the same earlier timestep s.The paths are evaluated in fully denoised action space before computing their difference.
  • The teacher generates xu from xt, providing the distillation signal used to train the student.
  • The final Consistency Policy objective combines CTM and DSM losses with tunable weights α and β.
  • Training uses stopgrad on teacher and selected student paths, leaving gradients through the t →s student generation.This prevents other operations from entering the computation graph and is intended to avoid unstable or failed convergence.
  • Experiments found adjacent t and u, with u = t −1 and arbitrary s below u, to work best.

C. Inference

Consistency Policy supports a speed–accuracy trade-off at inference without retraining. It offers single-step prediction for maximum speed and chained three-step generation for greater accuracy.

  • Both inference procedures remain faster than prior methods while allowing accuracy to be adjusted without further training.
  • Single-step inference samples z ∼N(0, I) and directly predicts the action with gθ(z, T, 0; o).The initial sample uses reduced variance relative to the standard N(0, T 2I) initialization.
  • Three-step inference repeatedly denoises, noises to a chaining timestep, and denoises again to refine the initial prediction.
  • Preset chaining timesteps avoid tuning separate steps for every task and dataset.
  • The three-step schedule prioritizes early-middle discretized timesteps using {t 2N 3 , t N 3 } for total step count N.

D. Implementation Details

The implementation preserves Diffusion Policy’s main infrastructure and input–output interface while evaluating simulation and real-world deployments against established baselines.

  • The experiments retain Diffusion Policy’s 1D Convolutional UNet architecture, image encoder, normalization, and action-sequence conventions.
  • The UNet conditions on observations and diffusion timestep t through FiLM blocks and processes actions with 1D convolutional blocks.
  • The architecture choice is treated as orthogonal to Consistency Policy, which could also use a properly tuned transformer backbone.
  • Real-world evaluation includes trash cleanup, plug insertion, and microwave interaction tasks.These tasks cover sequential cleanup, contact-rich precision, and long-horizon mobile-manipulator control.
  • The study compares methods using Diffusion Policy’s input and output formats, including two observation frames and end-effector-pose sequences.

B. Simulation Experiments

Simulation experiments cover six tasks across Robomimic, Push-T, and Franka Kitchen, measuring task performance alongside inference cost. Consistency Policy is substantially faster while retaining competitive accuracy, though long-horizon performance remains challenging.

  • Six tasks span image-based Robomimic control, state-based Push-T, and long-horizon Franka Kitchen evaluation.
  • Robomimic: Robomimic evaluates Lift, Can, Square, and Tool Hang with image observations and 200 proficient human demonstrations per task.
  • Franka Kitchen: Franka Kitchen tests multi-stage, long-horizon behavior, and single-step CP performs strongly early but struggles in later stages.
  • Metrics: Push-T reports target-area coverage, while Robomimic reports average success rate and standard error over 200 online evaluations.
  • Benchmark results: Single-step CP is at least an order of magnitude faster than DDPM and DDiM, while 3-step CP is competitive with DDPM accuracy.3-step CP is 3 and 9 times faster than DDiM and DDPM, respectively.
  • Benchmark results: On Robomimic Can, single-step CP marginally improves over DDPM and outperforms 3-step CP, whereas chaining helps much more on Tool Hang.
  • Inference speed: Simulation inference speeds are measured on an NVIDIA P5000 over 50 rollouts, with NFE dominating runtime.

C. Real World Experiments

Consistency Policy was evaluated on three real-world tasks using laptop-GPU inference, achieving substantially lower latency while preserving competitive task performance against a faster Diffusion Policy baseline.

  • Evaluation: Three real-world tasks—Trash Clean Up, Plug Insertion, and Microwave—were used to evaluate Consistency Policy.The tasks cover trash handling, contact-rich plug insertion, and long-horizon mobile manipulation.
  • Evaluation: The laptop setup used a single 3070 Ti GPU with 8GB of VRAM for inference on the first two tasks.The 100-step DDPM variant took around 1.5 seconds per forward pass, so DDiM was selected as the faster baseline.
  • Results: Consistency Policy achieved similar success rates to DDiM on Trash Clean Up and Plug Insertion while reducing latency by approximately 9x.The real-world comparison used average success rates and inference time measurements.
  • Results: In the Microwave task, Consistency Policy retained its inference-speed advantage but performed slightly worse than DDiM.The microwave experiment used a Kinova Gen3 arm mounted on a holonomic mobile base.
  • Results: Consistency Policy requires 15x fewer forward passes than DDiM, but overhead such as the observation encoder reduces the overall speedup to around 9x.The discrepancy is attributed to inference overhead outside the denoising passes.

D. Ablations

Ablations identify the consistency objective, initial sample variance, chaining schedule, teacher quality, and dropout as important design factors affecting training efficiency and policy performance.

  • Consistency Objective: The generalized CTM consistency objective performed best among three tested consistency objectives, while preserving the auxiliary DSM objective.The alternatives differed in the selected starting points and denoising stop point.
  • Consistency Objective: CTM and CTM-local had similar success rates and slightly outperformed Consistency Distillation, but CTM trained more than 40% slower.CTM’s higher training cost resulted from multiple teacher denoising steps between t and u.
  • Initial Sample Variance: Low-variance initial samples outperformed high-variance samples for both single-step and three-step Consistency Policy on Robomimic Square.The improvement was larger in the multi-step setting, potentially because intermediate noising preserves expressivity.
  • Preset Chaining Steps: Discretized subdivisions heavily outperformed continuous subdivisions on Tool Hang, while both achieved similar results on Square.The authors recommend trying subdivided discretized time first on difficult tasks.
  • Teacher Model Quality: Consistency Policy remained robust to teacher success rates across the tested range, with only a slight correlation between teacher quality and student success.The DSM objective is independent of the teacher and may help preserve student performance as teacher quality declines.
  • Dropout: Dropout contributed substantially to the consistency-objective training signal, and removing it from the s →0 region decreased success rate.The experiments used dropout set to 0.2.

V. LIMITATIONS

Consistency Policy improves raw inference speed while retaining performance, but trades away some multimodality and training stability and can require longer training.

  • Consistency Policy trades some of Diffusion Policy’s multimodality and stable training for faster inference.Its deterministic trajectory distillation can reduce multimodality, while the consistency objective makes training slightly less stable.
  • Consistency Policy slightly underperforms Diffusion Policy on accuracy in mobile manipulation experiments.
  • Consistency Policy generally requires more training time to reach Diffusion Policy’s performance, including teacher-training costs.Each training step runs the teacher model and multiple student-network forward passes.
  • Microwave distillation took longer than other tasks because Consistency Policy fits a harder single-step policy problem.The authors caution that these trade-offs should be evaluated for each use case.

VI. CONCLUSION

The paper presents consistency-based training for high-performing, low-latency visuomotor policies and evaluates it across simulation and real-world tasks.

  • Consistency Policy delivers substantially faster inference than prior diffusion-policy methods without sacrificing success rates across 9 simulation and real-world tasks.
  • Ablations identify consistency objective, lower initial sample variance, dropout, and preset chaining steps as important design choices.

APPENDIX A INFERENCE SPEEDS

The appendix defines end-to-end inference measurements and discusses speedups, sampling variance, experimental hardware, and possible explanations for variance-dependent behavior.

  • Inference Time includes image encoding, policy-network forward passes, and additional data-shuttling costs before actions reach the controller.
  • 13.3x network-time speedup contrasts with 9x end-to-end inference-time speedup between Consistency Policy and DDiM.
  • Actual practitioner speedups depend on the computation consumed by observation encoders and parallel processes alongside the policy network.Consistency Policy has greater impact when the policy network occupies a larger share of inference computation.
  • Lower-variance initial sampling improved generation quality over the higher-variance sampler used by EDM and Consistency Models.
  • Low-variance EDM sampling produced a gray block because the model lacked learned support in that region, although the result was heuristic.
  • The authors hypothesize that dimensionality differences and Gaussian mass concentration may explain weaker support for low-variance regions.They present this explanation as possible rather than certain.
  • Real-world stationary-arm experiments used a Franka Panda with two cameras, VR-collected demonstrations, and 15 Hz policy commands.
Loading 2405.07503v2…