MOSS: Mode-Oriented Source Shaping for Robot Flow Matching Policies

Zirun Ma, Jindou Jia, Tuo An, Jiaqi Bai, Jiahua Dong, Jianfei Yang · N/A · 2026

Summary

Abstract Summary

In robotic manipulation, the same observation often admits multiple feasible actions. Generative policies such as diffusion and flow matching naturally capture such multimodal behaviors, yet which mode is executed is left to random sampling and cannot be actively controlled. We propose Mode-Oriented Source Shaping (MOSS), which turns mode selection into an explicit choice of where generation starts. We first reveal, both empirically and theoretically, that deterministic flow transport induces mode-dependent structure in the source space. Under ideal deterministic ODE transport, distinct behavioral modes have distinct source-space preimages, while practical approximation errors can blur their boundaries due to imbalanced data, limited model capacity, and few inference steps. Motivated by this source-space structure, MOSS constructs semantically aligned mode-specific Gaussian sources that encode mode identity in the initial noise. A modality-strength estimator predicts the local degree of multimodality, which adaptively controls the geometry of the mode-specific sources. It contracts them toward a shared Gaussian where behaviors are similar and separates them near behavioral branching points. During deployment, the desired behavior is selected at inference simply by sampling from the corresponding component. Across 6 simulated and 3 real-world manipulation tasks, MOSS steers to the target mode up to 100% of the time with high task success, recovers minority modes at over 90% guidance rate under a 10:90 imbalance dataset, and surpasses mode-conditioned flow policy by 43.3 points of guidance rate on real robots.

Key Points

In robotic manipulation, the same observation often admits multiple feasible actions
Generative policies such as diffusion and flow matching naturally capture such multimodal behaviors, yet which mode is executed is left to random sampling and cannot be actively controlled.
We propose Mode-Oriented Source Shaping (MOSS), which turns mode selection into an explicit choice of where generation starts.
We first reveal, both empirically and theoretically, that deterministic flow transport induces mode-dependent structure in the source space.
Under ideal deterministic ODE transport, distinct behavioral modes have distinct source-space preimages, while practical approximation errors can blur their boundaries due to imbalanced data, limited model capacity, and few inference steps.

Multimodal Property of Flow Matching

Standard FM learns a continuous transport from a simple source to a complex action distribution conditioned on an observation $\mathbf{o}\in\mathcal{O}$ ( Lipman et al., 2023 ; Jia et al., 2026b ) .

Let $d\in\mathbb{N}$ denote the dimension of the vectorized action chunk, $\mathbf{z}\in\mathbb{R}^{d}$ with $\mathbf{z}\sim\mu_{0}=\mathcal{N}(\mathbf{0}_{d},\mathbf{I}_{d})$ be the source noise, where $\mathbf{0}_{d}\in\mathbb{R}^{d}$ and $\mathbf{I}_{d}\in\mathbb{R}^{d\times d}$, and let $\mathbf{y}\in\mathbb{R}^{d}$ with $\mathbf{y}\sim p^{\math...

Mode-Oriented Source Shaping

Section 3 shows that FM naturally induces mode-dependent regions in the source space, but this separation is implicit and can become weak or entangled in practice. Rather than recovering these implicit regions, MOSS deliberately imposes a simple, semantically aligned source geometry to make mode separation explicit and controllable.

It enlarges the source-space separation between modes when it is actually needed, and selects the mode by choosing where the transport starts .

Mode-Oriented Source Shaping — MOSS: Mode-Oriented Source Shaping for Robot Flow Matching Policies — Figure 2: Initial-noise mode separation in FM. Experimental details are provided in Appendix A .

Evaluation

We evaluate MOSS on $6$ simulated tasks and $3$ real-world robotic tasks, as shown in Figure 3 . The simulation benchmark consists of $4$ bimodal tasks— Push Cube , Grasp Eyeglasses , Stack Cube , and Put Wine —and two multi-mode tasks, 2D Navigation and Place Cube ( Jia et al., 2026a ; Mu et al., 2021 ; Liu et al., 2023 ) . Each task admits multiple distinct yet valid strategies for successful completion.

For real-world evaluation, we consider Place Cups , Pick Veg , and Open Drawer on a Galaxea R1 Lite robotic platform.

Evaluation — MOSS: Mode-Oriented Source Shaping for Robot Flow Matching Policies — Figure 3: Overview of the simulation and real-world tasks used for evaluation.

Sources and Demonstrations

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