MOSS: Mode-Oriented Source Shaping for Robot Flow Matching Policies
MOSS, a preprint from Nanyang Technological University's MARS Lab and Tsinghua University, lets an operator choose which of several valid strategies a flow-matching robot policy executes by picking where its noise sampling starts, and on a Galaxea R1 Lite robot it hit the requested behaviour 84.2% of the time on average against 40.8% for a mode-conditioned policy, a 43.3-point gap. Average task success on the three real tasks rose to 85.3% from 73.3%.
Generative robot policies can learn several valid ways to do the same task, such as approaching from the left or the right. Can a user or planner choose which one the robot executes, reliably, even when the training data contain far more examples of one way than the other?
Diffusion and flow-matching policies represent multiple behaviours, but which one runs is left to random sampling. The usual fix, feeding the desired mode as a condition such as a language instruction, often fails when demonstrations are imbalanced: the network learns shortcuts from the rich visual scene to the action and ignores the weak, low-dimensional condition. Learned Gaussian-mixture priors in the source space carry no explicit meaning and need extra networks to align with user-defined modes.
Mode-conditioned flow matching passes the requested mode to the velocity network; GMM-sourced flow matching learns mixture components without semantic labels; unconditional policies such as Diffusion Policy, ACT and BET reproduce multiple behaviours but give no control over which appears.
The authors show that with deterministic ODE transport, different behaviour modes already come from distinct, though entangled, regions of the starting noise. MOSS makes that structure explicit: each mode gets its own Gaussian source, and a modality-strength estimator measures how much the possible futures diverge at the current state, pulling the sources together where behaviours look alike and pushing them apart near branching points. The mode label only selects the starting distribution and is never fed to the shared velocity field, so choosing a behaviour at deployment means sampling from the matching component.
On three real tasks on a Galaxea R1 Lite (place cups, pick vegetables, open a drawer), MOSS reached guidance rates of 0.95 and 1.00, 0.80 and 0.65, and 0.90 and 0.75 for the two modes, while the mode-conditioned baseline scored 0.00 for one mode on both the vegetable and drawer tasks; MEKAWIRE's average of the six rates is 84.2% against 40.8%, consistent with the paper's 43.3-point gap. Average real success was 85.3% against 73.3% for the mode-conditioned policy and 68.3% for plain FM-DiT, though the baseline was slightly better on the drawer (0.85 vs 0.83). In four bimodal simulation tasks the average guidance rate was 96.4% against 94.5% for mode conditioning and 74.0% for GMM sources; MOSS trailed on one mode of Put Wine (0.86 vs 0.94). With a 10:90 imbalanced dataset it still steered to the minority mode more than 90% of the time. Ablations on three tasks: one shared Gaussian averaged 67.0% guidance, fixed separation 85.5%, full MOSS 98.8%.
This is a preprint. The modes must be a fixed, discrete set defined before training, so open-ended or continuous behaviour control is not covered, and the authors say performance inside large generative robot policies is still to be tested. The real-robot study covers three two-mode tasks on one platform. In simulation the advantage over plain mode conditioning is small (about two points on average), so most of the benefit shows up on real hardware and with imbalanced data, where results rest on a limited number of tasks.
Flow matching is the action head of Physical Intelligence's pi0 family and of many humanoid and dual-arm stacks, so a cheap way to choose a behaviour at inference without retraining matters to anyone deploying them with operator preferences, safety constraints or planner commands, for example a fixed approach side in a crowded cell. Galaxea AI's R1 Lite is the test platform. Time scale: easy to try now in research stacks; proof inside large VLA models is the next step and probably a year or more out.
The full text is not republished here because the paper's license does not allow it. Read the original on arXiv.