Imitation Learning via Discrete and Rhythmic Primitive Composition

Abstract

Long-horizon robotic manipulation is inherently composed of semantically meaningful sub-tasks. While Movement Primitives (MPs) can model these components, their composition typically relies on manual engineering or sample-inefficient reinforcement learning. Existing supervised methods are limited by the lack of MP labels in standard datasets and often fail to generalize across discrete and rhythmic movements. We propose Neural Dynamic Composition (NDC), a supervised framework to compose Dynamic Movement Primitives (DMPs) using only raw trajectories for supervision, given a single example of each primitive. Our approach exploits parallel combinations of DMPs as a differentiable relaxation of categorical primitive selection, and utilizes a warp-invariant loss formulation to identify primitives despite spatial or temporal variations. A staged training curriculum is employed to learn the selection policy and DMP parameters, unifying discrete and rhythmic movements in one framework. Evaluated on 2D handwriting and 3D robotic tasks, our method outperforms baselines on most tasks and metrics. An ablation study demonstrates how each design choice contributes to the system's performance.