Robot Learning · Flow Matching · Multimodal Control

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

Controllable robot behavior starts at the source.

Zirun Ma1,2,*Jindou Jia1,*Tuo An1 Jiaqi Bai1Jiahua Dong1Jianfei Yang1,†
1 MARS Lab, Nanyang Technological University 2 Tsinghua University

* Equal contribution  ·  † Corresponding author

Paper Coming soon arXiv Coming soon Code Coming soon
MOSS overview showing a multimodal robot task, adaptive source separation, and the shared flow network
One policy, multiple behaviors. MOSS adaptively shapes the source distribution so that choosing where generation starts determines the requested mode.

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.

up to 100%target-mode guidance with high task success
>90% GRfor minority modes under a 10:90 data imbalance
+43.3 ppaverage real-world GR over mode-conditioned FM

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

Place Cups · Mode 1Placeholder
Place Cups · Mode 2Placeholder

Pick Veg

Pick Veg · Mode 1Placeholder
Pick Veg · Mode 2Placeholder

Open Drawer

Open Drawer · Mode 1Placeholder
Open Drawer · Mode 2Placeholder

Simulation Overview

Simulation montagePlaceholder

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.

MOSS makes this implicit organization explicit and semantically aligned.
Overlapping clusters of initial noise corresponding to four behavioral modes
Initial-noise mode separation in a standard flow-matching policy.

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.

01

Estimate ambiguity

A lightweight estimator predicts the local modality strength from observation and action history.

02

Shape the source

Mode-specific Gaussian sources merge along shared motion and separate as future behaviors diverge.

03

Select at inference

A user or planner chooses mode k; sampling its source component selects the requested behavior.

One shared flow networkNo mode label injected into the velocity fieldNatural extension to multiple modes

Experiments

Six simulated and three real-world tasks.

Each task admits multiple distinct yet valid strategies for successful completion.

Six simulation tasks and three real-world robot tasks used to evaluate MOSS
Simulation benchmarks and real-world tasks on a Galaxea R1 Lite robot.

Main results

Reliable control without sacrificing success.

GR measures whether a rollout both succeeds and realizes the requested semantic mode; SR measures task success.

01

Controllable behavior in simulation

MOSS maintains high guidance and success across four bimodal tasks, while source components remain explicitly aligned with semantic modes.

Guidance rate / success rate (GR / SR)
MethodPush CubeGrasp EyeglassesStack CubePut Wine
M1M2M1M2M1M2M1M2
Mode-Cond. FM.86 / .96.94 / 1.001.00 / 1.001.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 / .991.00 / 1.001.00 / 1.001.00 / 1.00.98 / .98.96 / .96.92 / .94.86 / .94
Behavior distributions of unconditional policies across four simulation tasks
Unconditional policies can represent multiple outcomes, but cannot select which one to execute.
02

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.

Average success rate and mode-specific guidance rates
TaskMethodSRGR₁GR₂
Place CupsFM-DiT.70——
Mode-Cond. FM.70.15.85
MOSS.98.951.00
Pick VegFM-DiT.60——
Mode-Cond. FM.65.00.60
MOSS.75.80.65
Open DrawerFM-DiT.75——
Mode-Cond. FM.85.85.00
MOSS.83.90.75
03

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.

Minority-mode learning curves under imbalanced training datasets
Minority-mode guidance rate versus exposure to minority-mode training samples.
04

Why adaptation matters

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

Guidance-rate curves comparing adaptive MOSS with fixed source shaping
Checkpoint-wise guidance rates for adaptive MOSS and fixed full-chunk source shaping.

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}
}