Robot Learning · Flow Matching · Multimodal Control
MOSS: Mode-Oriented Source Shaping
for Robot Flow Matching Policies
Controllable robot behavior starts at the source.
* Equal contribution · † Corresponding author
Abstract
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.
MOSS in action
Video Demonstrations
Different modes. One shared policy. Explore the real-world manipulation tasks and simulation setting.
Preview footage · All players currently use the same placeholder clip. Task-specific rollouts will replace these previews.
Place Cups
Pick Veg
Open Drawer
Simulation Overview
Key insight
Flow matching already organizes its noise by behavior.
Under deterministic transport, each behavioral mode has a distinct preimage in the initial noise space. We call these regions noise zones.
In practice, finite data, limited model capacity, and few integration steps blur their boundaries. The plot shows that the initial noise is predictive of the resulting mode, but neighboring modes still overlap too much for reliable control.

Method
Control the mode through the source—not the network.
MOSS preserves a single shared velocity field and changes only the initial distribution used for generation.
Estimate ambiguity
A lightweight estimator predicts the local modality strength from observation and action history.
Shape the source
Mode-specific Gaussian sources merge along shared motion and separate as future behaviors diverge.
Select at inference
A user or planner chooses mode k; sampling its source component selects the requested behavior.
Experiments
Six simulated and three real-world tasks.
Each task admits multiple distinct yet valid strategies for successful completion.

Main results
Reliable control without sacrificing success.
GR measures whether a rollout both succeeds and realizes the requested semantic mode; SR measures task success.
Controllable behavior in simulation
MOSS maintains high guidance and success across four bimodal tasks, while source components remain explicitly aligned with semantic modes.
| Method | Push Cube | Grasp Eyeglasses | Stack Cube | Put Wine | ||||
|---|---|---|---|---|---|---|---|---|
| M1 | M2 | M1 | M2 | M1 | M2 | M1 | M2 | |
| Mode-Cond. FM | .86 / .96 | .94 / 1.00 | 1.00 / 1.00 | 1.00 / 1.00 | .96 / .96 | .94 / .94 | .92 / .92 | .94 / .96 |
| GMM-Sourced FM | .72 / .96 | .92 / 1.00 | .98 / 1.00 | .42 / 1.00 | .58 / 1.00 | .56 / 1.00 | .94 / .94 | .80 / .94 |
| MOSS | .99 / .99 | 1.00 / 1.00 | 1.00 / 1.00 | 1.00 / 1.00 | .98 / .98 | .96 / .96 | .92 / .94 | .86 / .94 |

Balanced control on real robots
Direct conditioning collapses on one requested mode in Pick Veg and Open Drawer; MOSS retains control of both modes under visual and execution variability.
| Task | Method | SR | GR₁ | GR₂ |
|---|---|---|---|---|
| Place Cups | FM-DiT | .70 | — | — |
| Mode-Cond. FM | .70 | .15 | .85 | |
| MOSS | .98 | .95 | 1.00 | |
| Pick Veg | FM-DiT | .60 | — | — |
| Mode-Cond. FM | .65 | .00 | .60 | |
| MOSS | .75 | .80 | .65 | |
| Open Drawer | FM-DiT | .75 | — | — |
| Mode-Cond. FM | .85 | .85 | .00 | |
| MOSS | .83 | .90 | .75 |
Minority-mode efficiency
With only 10 minority-mode demonstrations against 90 majority-mode demonstrations, MOSS rapidly recovers the requested minority behavior by reusing shared motion structure.

Why adaptation matters
Full-chunk, adaptive source shaping reaches consistently high guidance rates and remains more stable than fixed separation at later checkpoints.

Citation
BibTeX
@misc{ma2027moss,
title = {MOSS: Mode-Oriented Source Shaping for Robot Flow Matching Policies},
author = {Ma, Zirun and Jia, Jindou and An, Tuo and Bai, Jiaqi and Dong, Jiahua and Yang, Jianfei},
year = {2027}
}
