Fast Non-Parametric Heteroscedastic Imitation Learning With Geometric Priors

Maximilian Mühlbauer; Arne Sachtler; Markus Knauer; Cem Küçükgenç; Yanlong Huang; Alin Albu-Schäffer; João Silvério

cs.ROOct 6, 2026

Fast Non-Parametric Heteroscedastic Imitation Learning With Geometric Priors

Authors: Maximilian Mühlbauer, Arne Sachtler, Markus Knauer, Cem Küçükgenç, Yanlong Huang, Alin Albu-Schäffer, João Silvério

Organizations: Department of Computer Engineering, Technical University of Munich, Friedrich-Ludwig-Bauer-Str. 3, Garching, Germany. · German Aerospace Center (DLR), Robotics and Mechatronics Center (RMC), Münchener Str. 20, 82234 Weßling, Germany. · School of Computer Science, University of Leeds, Leeds, UK.

Abstract

When learning probabilistic policies from human demonstrations, data-efficient learning and fast adaptations to new scenarios are key requirements. One popular way to achieve intuitive and reliable adaptations is through non-parametric, typically kernel-based, methods. However, existing solutions either fail to account for the geometry of manifolds common in robotics, limiting data efficiency, or, when geometry-aware, provide unreliable uncertainty estimates or require retraining to adapt. We propose a non-parametric approach leveraging geometric priors in scenarios of data scarcity and heteroscedastic uncertainties for probabilistic modeling. We utilize the method to formulate policies based on time or robot state, where non-separable diagonal kernels allow capturing uncertainty relations between degrees of freedom for same-sized in- and outputs. Fast updates, requiring less than 3 ms for a trajectory involving both position and orientation are possible through an optimized formulation. Our approach supports both manifold-valued input and manifold-valued output with large orientation changes. Using task parameterization, adaptation to different object poses is easily possible. We evaluate the approach on a set of toy examples and on real robot manipulation tasks both in autonomous execution and in shared control scenarios.

Figures & tables

Fig. 2 : Manifolds used in this work, proposed by [ 21 ] . Cylindrical (

\mathcal{M}_{2}

) or spherical (

\mathcal{M}_{3}

) coordinates allow expressing some tasks more data-efficient compared to Cartesian (

\mathcal{M}_{1}

) coordinates. Coordinate systems in each image depict the unit quaternion

(0,0,0,1)^{\top}

for different positions on the manifold.

Fig. 3 : Left : Pose data 1 reproduced by a Riemannian KMP ( III-A ) 2 with correct rotations. In contrast, the STS KMP [ 9 ] 3 suffers from incorrect rotations. Right : Adding a via point V

0.1\text{\,}\mathrm{m}

above the trajectory with a relative rotation of

-0.35\text{\,}\mathrm{rad}

about the

y

axis and a null space action N . Via point and nullspace target are visualized by a larger coordinate system.

Fig. 4 : Plots for the letter “B” of the handwriting dataset [ 10 ] projected to

\mathcal{S}^{2}

. We fit the GMM using

10

Gaussians and use

20

reference points for KMP ; all these Gaussians are depicted with their mean and covariance ellipsoid.

Method	Position Difference (mm)	Rotation Difference (mrad)
GMR ( STS )	$2.47\pm 1.02$	$3139.36\pm 11.47$
KMP ( STS ) [ 9 ]	$2.47\pm 1.03$	$269.36\pm 646.48$
Riemannian GMR [ 4 ]	$2.47\pm 1.09$	$15.27\pm 8.81$
Riemannian KMP (ours)	$\bm{2.46\pm 1.12}$	$\bm{15.15\pm 9.05}$
Nadaraya-Watson [ 17 ]	$3.05\pm 1.22$	$15.29\pm 9.15$

TABLE I : Differences between the prediction of different methods and ground truth for the spiral trajectory ( Fig. 3 ).

Kernel	$\mathcal{S}^{1}$	$\mathcal{S}^{1}\times\mathbb{R}^{2}$	$\mathcal{S}^{2}$	$\mathrm{SO}(3)$	$\mathbb{R}^{3}\times\mathrm{SO}(3)$
RBF	$0.01$	$0.65$	$0.27$	$0.63$	$1.1$
Matérn ( $\nu=\frac{3}{2}$ )	$0.40$	$0.99$	$0.77$	$0.97$	$1.46$
Matérn ( $\nu=\frac{5}{2}$ )	$0.07$	$0.83$	$0.58$	$0.8$	$1.33$

TABLE II : Maximum length scales for selected kernels.

Method	Fit ( V-A )	Add ( V-B )	Remove ( V-C )	Predict ( V-A )
N.-W. [ 17 ]	$0.12\pm 0.00$	-/-	-/-	$0.83\pm 0.72$
KMP [ 8 ]	$12.93\pm 3.93$	$13.63\pm 3.97$	$11.42\pm 3.16$	$1.68\pm 0.33$
iKMP [ 20 ]	$12.17\pm 3.01$	$1.18\pm 0.09$	$0.84\pm 2.33$	$1.71\pm 0.58$
cKMP ( V )	$\bm{2.66\pm 0.78}$	$\bm{0.82\pm 0.16}$	$\bm{0.60\pm 0.58}$	$\bm{1.35\pm 0.59}$

TABLE III : Computation times (milliseconds) for adding and removing

3

points and predicting

200

points for letter writing.

Method	Fit ( V-A )	Add ( V-B )	Remove ( V-C )	Predict ( V-A )
N.-W. [ 17 ]	$0.22\pm 0.01$	-/-	-/-	$24.06\pm 2.89$
R. KMP ( III )	$37.11\pm 3.74$	$37.20\pm 1.72$	$34.51\pm 2.61$	$13.28\pm 1.22$
iKMP [ 20 ]	$36.22\pm 3.35$	$4.11\pm 1.60$	$2.29\pm 0.52$	$13.15\pm 0.95$
cKMP ( V )	$\bm{7.89\pm 2.32}$	$\bm{2.64\pm 0.32}$	$\bm{1.27\pm 0.37}$	$\bm{8.57\pm 0.82}$

TABLE IV : Computation times (milliseconds) fitting, adding

3

poses and subsequently removing them in a trajectory with

100

reference poses, predicting

500

poses.

Explore similar work

Dynamics Are Learned, Not Told: Semi-Supervised Discovery of Latent Dynamics Geometries For Zero-Shot Policy Adaptation
Jun 1, 2026Zhiming Xu, Weitao Zhou, Xianghui Pan +4Discriminative Congruence Transform Zero-Shot
School of Computer Science and Engineering, Tongji University, Shanghai, China · School of Vehicle and Mobility, Tsinghua University, Beijing, China · Simple AI, Beijing, China +1
Improving Imitation Learning Efficiency for Manipulation through Geometric Prior Pretraining
Sep 14, 2026Shogo Iwakata, Tomohiro Motoda, Ryosuke Yamada +8Temporally Coherent Imitation Learning Imitation Learning
Waseda University · Embodied AI Research Team, AIRC and National Institute of Advanced Industrial Science and Technology (AIST) · Computer Vision Research Team, AIRC and National Institute of Advanced Industrial Science and Technology (AIST) +3
Proxy Policy Steering
Sep 8, 2026Chuanruo Ning, Tianrui Wang, Wei-Chiu Ma +1Robot Policies Robot Systems
Cornell University

Jun 1, 2026, cs.ROJ/K move · Enter open · S save

Dynamics Are Learned, Not Told: Semi-Supervised Discovery of Latent Dynamics Geometries For Zero-Shot Policy Adaptation

Zhiming Xu, Weitao Zhou, Xianghui Pan, Nanshan Deng+3

School of Computer Science and Engineering, Tongji University, Shanghai, China · School of Vehicle and Mobility, Tsinghua University, Beijing, China · Simple AI, Beijing, China+1

Discriminative Congruence Transform Zero-Shot Discovery

Open paper AskPDF arXiv

Real-world dynamics shifts pose a critical challenge for reinforcement learning in robotics, as policies tightly coupled to nominal environments often fail catastrophically when physical conditions change. Most existing methods rely on encoding explicitly identified physical parameters into a latent context, a parameter-centric paradigm that depends on pre-specified axes of variation and becomes brittle under unmodeled or compound dynamics changes. We revisit dynamics adaptation from an outcome-centric perspective: rather than telling policies what the dynamics are, we enable them to learn how dynamics affect interaction outcomes. Theoretically, this is grounded in a monotonic relationship between target-domain regret and the Lipschitz constant of a trajectory dynamics encoder. Practically, this constant can be upper-bounded through contrastive learning, yielding a smooth, task-relevant latent topology without privileged dynamics information. On MuJoCo benchmarks, our method consistently outperforms parameter-centric baselines under severe dynamics shifts, including unmodeled and time-varying parameters, while also improving in-distribution stability and latent interpretability. Overall, these results validate that controlling latent geometry is a principled mechanism for robust adaptation.

Abstract

Explore similar work

Dynamics Are Learned, Not Told: Semi-Supervised Discovery of Latent Dynamics Geometries For Zero-Shot Policy Adaptation

Improving Imitation Learning Efficiency for Manipulation through Geometric Prior Pretraining

Proxy Policy Steering