Kinematics-Aware Diffusion Policy with Consistent 3D Observation and Action Space for Whole-Arm Robotic Manipulation

TL;DR

KADP introduces a 3D node-based representation for whole-arm robotic manipulation, aligning task, observation, and action spaces to improve sample efficiency and spatial generalization. By embedding kinematic constraints directly into a diffusion-policy framework and solving for joint trajectories via a whole-body IK solver, the method enables feasible, full-arm control from limited demonstrations. Extensive RLBench simulations and real-world Panda experiments show superior performance and robustness over end-effector and joint-space baselines, especially in collision-avoidant and body-interaction scenarios. The work highlights the value of consistent 3D representations and kinematic priors for scalable, body-aware manipulation learning.

Abstract

Whole-body control of robotic manipulators with awareness of full-arm kinematics is crucial for many manipulation scenarios involving body collision avoidance or body-object interactions, which makes it insufficient to consider only the end-effector poses in policy learning. The typical approach for whole-arm manipulation is to learn actions in the robot's joint space. However, the unalignment between the joint space and actual task space (i.e., 3D space) increases the complexity of policy learning, as generalization in task space requires the policy to intrinsically understand the non-linear arm kinematics, which is difficult to learn from limited demonstrations. To address this issue, this letter proposes a kinematics-aware imitation learning framework with consistent task, observation, and action spaces, all represented in the same 3D space. Specifically, we represent both robot states and actions using a set of 3D points on the arm body, naturally aligned with the 3D point cloud observations. This spatially consistent representation improves the policy's sample efficiency and spatial generalizability while enabling full-body control. Built upon the diffusion policy, we further incorporate kinematics priors into the diffusion processes to guarantee the kinematic feasibility of output actions. The joint angle commands are finally calculated through an optimization-based whole-body inverse kinematics solver for execution. Simulation and real-world experimental results demonstrate higher success rates and stronger spatial generalizability of our approach compared to existing methods in body-aware manipulation policy learning.

Figures (7)

• Figure 1: The proposed approach uses a set of 3D nodes on the arm body as both robot state and action representation for whole-arm manipulation, which is consistent with the 3D point cloud observation space and task space. Compared with using end-effector poses or joint angles, our method achieves higher spatial generalizability and sample efficiency while ensuring kinematic feasibility.
• Figure 2: Overview of Kinematics-Aware Diffusion Policy (KADP). Taking the encoded 3D visual representations, the 3D robot nodes and time embeddings as input, diffusion model predicts the denoised 3D node trajectory iteratively. For execution, the joint angle commands are computed through an optimization-based whole-body inverse kinematics solver.
• Figure 3: Visualization of 8 RLBench simulation tasks.
• Figure 4: Ablation on the kinematic constraints in DP and the number of nodes. IK Error refers to the average per-node inverse kinematics optimization error when solving joint commands. KADP w/o KC: remove the kinematic constraints in DP. Node-3 / Node-5: replace the full 8 nodes with fewer nodes.
• Figure 5: Overview of the 4 real-world tasks, where the manipulation is achieved by the proposed KADP.