Adaptive Diffusion Constrained Sampling for Bimanual Robot Manipulation

TL;DR

ADCS presents a diffusion-based framework for constraint-aware sampling in high-dimensional multi-DoF robotic systems, combining per-constraint energy networks with a Transformer-based Compositional Weighting Transformer to adaptively balance equality and inequality constraints in $SE(3)$. It introduces a two-stage, batch-wise sampling strategy that couples Annealed Langevin dynamics with Gauss-Newton refinement and KDE-based resampling to achieve diverse, high-quality samples. The method demonstrates improved constraint satisfaction, coverage, and generalization in both simulated dual-arm manipulation and real-world stippling tasks on Franka and TIAGo robots, outperforming baselines like Gauss-Newton, Diffusion-CCSP, and NLP-Sampling. The work contributes a modular, adaptable approach for constraint fusion in diffusion models, with potential as a planning primitive and for future integration with vision-language guidance.

Abstract

Coordinated multi-arm manipulation requires satisfying multiple simultaneous geometric constraints across high-dimensional configuration spaces, which poses a significant challenge for traditional planning and control methods. In this work, we propose Adaptive Diffusion Constrained Sampling (ADCS), a generative framework that flexibly integrates both equality (e.g., relative and absolute pose constraints) and structured inequality constraints (e.g., proximity to object surfaces) into an energy-based diffusion model. Equality constraints are modeled using dedicated energy networks trained on pose differences in Lie algebra space, while inequality constraints are represented via Signed Distance Functions (SDFs) and encoded into learned constraint embeddings, allowing the model to reason about complex spatial regions. A key innovation of our method is a Transformer-based architecture that learns to weight constraint-specific energy functions at inference time, enabling flexible and context-aware constraint integration. Moreover, we adopt a two-phase sampling strategy that improves precision and sample diversity by combining Langevin dynamics with resampling and density-aware re-weighting. Experimental results on dual-arm manipulation tasks show that ADCS significantly improves sample diversity and generalization across settings demanding precise coordination and adaptive constraint handling.

Figures (7)

• Figure 1: Overview of the Adaptive Diffusion Constrained Sampling (ADCS) architecture. The model consists of three components: a constraint feature encoder, an energy network (MLP), and the Compositional Weighting Transformer (CWT). During training, noisy poses in $\mathrm{SE}(3)^n$ are used to learn constraint-aware energy functions. In inference, sampling is performed in batched joint space using Langevin dynamics guided by the learned energy landscape.
• Figure 2: Sampled configurations for simulation tasks.
• Figure 3: The real-world experimental setup involves two Franka robots collaboratively grasping a pen for stippling operations and a TIAGo to carry objects using both hands.
• Figure 4: Final stippling layout of 8 different real-world tasks.
• Figure 5: This figure illustrates the one end-effector's position distribution densities along the x, y, and z axes before replication, after uniform replication, and after Gaussian-kernel-based replication.