Contractive Dynamical Imitation Policies for Efficient Out-of-Sample Recovery

TL;DR

The paper addresses the brittleness of imitation learning in out-of-sample scenarios by introducing contractive dynamical-system policies (SCDS) built from recurrent equilibrium networks (RENs) and coupling layers to guarantee contraction for any parameter values. By solving a differentiable IVP via Neural ODEs, and employing a trajectory-space loss with differentiable soft-DTW, the method learns state-only policies that robustly recover from perturbations and unseen initial states. Theoretical guarantees are provided: an upper bound on the deployment loss decomposes into a data-fit term and a contraction-dependent term, with a corollary bounding the true loss under an initial-condition region. Empirically, SCDS demonstrates superior out-of-sample recovery on LASA and Robomimic datasets, outperforming baselines like SNDS, SDS-EF, and BC, and enabling practical deployment in robotic simulators, thereby offering reliable, contractive imitation in high-dimensional control tasks.

Abstract

Imitation learning is a data-driven approach to learning policies from expert behavior, but it is prone to unreliable outcomes in out-of-sample (OOS) regions. While previous research relying on stable dynamical systems guarantees convergence to a desired state, it often overlooks transient behavior. We propose a framework for learning policies modeled by contractive dynamical systems, ensuring that all policy rollouts converge regardless of perturbations, and in turn, enable efficient OOS recovery. By leveraging recurrent equilibrium networks and coupling layers, the policy structure guarantees contractivity for any parameter choice, which facilitates unconstrained optimization. We also provide theoretical upper bounds for worst-case and expected loss to rigorously establish the reliability of our method in deployment. Empirically, we demonstrate substantial OOS performance improvements for simulated robotic manipulation and navigation tasks.

Figures (21)

• Figure 1: Policy rollouts generated by contractive and stable policies. While both policies eventually reach the target, the contractive policy closely mimics the expert in the transient phase.
• Figure 2: Overview of the SCDS training scheme. The policy structure (top box) consists of a REN, a linear projection, and a bijection block, ensuring contractivity for any choice of parameters, $\boldsymbol{\theta}$. In the forward pass (→), initial states are passed through a differentiable ODE solver to generate state rollouts. The loss function penalizes the discrepancy between the generated and expert trajectories, updating the policy via backpropagation (←).
• Figure 3: Uniform weighting averages the deviation from both demonstrations, hence, $L (\mathbf{y}_0^a; \; \boldsymbol{\theta}) \neq 0$. In contrast, \ref{['eq:lambda']} results in $L (\mathbf{y}_0^a; \; \boldsymbol{\theta}) = 0$. For $\mathbf{y}_0^c$, the weights in \ref{['eq:lambda']} assign higher importance to the closer $\mathbf{y}_0^b$'s demonstration, which is intuitive.
• Figure 4: In-sample and OOS policy rollouts for selected 2D tasks in the LASA dataset. The training process promotes higher contraction rates ( \ref{['app:learning_contraction_rate']}), resulting in effective OOS recovery.
• Figure 5: Comparing SCDS to the selected baselines on in-sample and OOS rollouts in the 2D task space. BC results in diverging trajectories due to the lack of stability guarantees. While SNDS and SDS-EF ensure global stability and reach the target, they display large deviations from expert data. In contrast, SCDS efficiently recovers from OOS states through its contracting transient behavior.

Theorems & Definitions (14)

• Definition 2.1
• Proposition 2.1
• Remark 3.1
• Theorem 4.1
• Corollary 4.1.1
• Lemma A.1
• proof
• proof
• Lemma A.2
• proof : Proof of Lemma \ref{['lemma:invinvineq']}