Faster Visuomotor Policy Learning on Action Manifolds via Riemannian MeanFlow

S. Talha Bukhari, Austin Garrett, Yi Wei, Ruiqi Ni, Zachary Kingston, Aniket Bera · N/A · 2026

Summary

Abstract Summary

Visuomotor policies learn a direct map from raw sensory observations to robot action sequences. Policies based on Diffusion and Flow Matching capture the multimodal distribution over action sequences in an end-to-end manner. This expressivity comes at the cost of multi-step numerical integration of the learned vector field for action generation, which can be expensive and time-consuming, impeding fast control rates required in robotics applications. Furthermore, robot action sequences are usually defined on a smooth, differentiable manifold, requiring that the learned policy respects the intrinsic geometry of the robot's action space. Here, we present Riemannian MeanFlow Policy (RMFP), which learns the conditioned flow map of the probability path on the robot action manifold. Our formulation employs a flow map consistency objective grounded in the data by a Riemannian Conditional Flow Matching anchor. The flow map consistency condition is stable to train and constrains the learned model to finite-time transport, which yields on-manifold action sequence generation with as few as one network function evaluation. We present results on the spherical LASA and Push-T benchmarks, on the Tool Hang and Transport tasks of the Robomimic suite, and on the Franka Kitchen task with manifold-constrained action generation, and demonstrate that RMFP attains performance competitive with prior work at a lower sampling cost. We also employ RMFP on a real-world robotic manipulation task to demonstrate fast action generation under imperfect sensor measurements in the physical world.

Key Points

Visuomotor policies learn a direct map from raw sensory observations to robot action sequences
Policies based on Diffusion and Flow Matching capture the multimodal distribution over action sequences in an end-to-end manner.
This expressivity comes at the cost of multi-step numerical integration of the learned vector field for action generation, which can be expensive and time-consuming, impeding fast control rates required in robotics applications.
Furthermore, robot action sequences are usually defined on a smooth, differentiable manifold, requiring that the learned policy respects the intrinsic geometry of the robot's action space.
Here, we present Riemannian MeanFlow Policy (RMFP), which learns the conditioned flow map of the probability path on the robot action manifold.

Methodology

The action space of a manipulator is a smooth manifold rather than a vector space. Let $\mathcal{A}$ denote the manifold of a single action, which for an end-effector controller is the product $\mathcal{A}=\mathbb{R}^{3}\times\mathcal{S}^{3}\times\mathbb{R}^{2}$ of a position, a unit quaternion and two finger positions [ 2 ] .

Experiments

We ask whether a learned action flow map on a Riemannian manifold retains task performance at a limited sampling budget under different settings. We evaluate the proposed formulation against baselines for Riemannian generative policy learning on four simulated benchmarks and one trajectory-reproduction task. We use four simulated benchmarks (cf. Fig. 3 ) and one trajectory-reproduction task (cf. Fig.

1 ) that separate the effect of the training objective from the effect of the geometry of the action space.

Experiments — Faster Visuomotor Policy Learning on Action Manifolds via Riemannian MeanFlow — (b) Robomimic Tool Hang

Sources and Demonstrations

Figures are reproduced from Bukhari et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.