Source-linked AI summary
MOSS: Mode-Oriented Source Shaping for Robot Flow Matching Policies
Zirun Ma, Jindou Jia, Tuo An, Jiaqi Bai, Jiahua Dong, Jianfei Yang
TL;DR
Generative robot policies can represent multiple feasible behaviors, but sampling leaves the executed mode uncontrolled. MOSS encodes semantic modes in adaptive Gaussian sources for flow matching. Across simulated and real-world manipulation tasks, it achieves strong mode controllability and task success, including gains over mode-conditioned flow matching on real robots.
Problem
Generative policies can represent multiple feasible behaviors, but random sampling largely determines which mode executes, limiting user-specified control.
Source: §1 · p. 1Method
MOSS uses semantically defined mode-specific Gaussian sources and adapts their separation according to estimated local multimodality.
Source: §1 · p. 2Results
On real-world tasks, MOSS averages 84.2% GR/85.3% SR, outperforming mode-conditioned FM by 43.3/12.0 percentage points; it also reaches up to 100% target-mode steering across 6 simulated and 3 real-world tasks.
Source: §1 · p. 3Takeaways & Limitations
Source-space geometry is an effective design dimension for controllable generative robot policies, with reported gains in minority-mode learning under imbalanced demonstrations.
Source: §6 · p. 10MOSS assumes a fixed set of discrete modes specified before training and has not yet been evaluated in larger-scale generative robot policies.
Source: §7 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
Generative policies represent multiple feasible behaviors, but random sampling leaves mode choice implicit. MOSS analyzes the source-space structure of flow matching and shapes mode-specific Gaussian sources to make selection explicit and adaptive.
1 Introduction
- Because mode choice in existing generative policies largely follows random sampling, they offer limited control for user preferences, test-time constraints, or human intervention.Generative policies can produce multiple feasible behaviors for the same observation.Source: §1 · p. 1
- Mode-conditioned policies can fail under imbalanced training, while learned Gaussian priors lack semantic alignment and require extra networks and post-alignment.The source-prior limitation is described for existing source-based methods.Source: §1 · pp. 1–2
- Ideal deterministic ODE transport partitions source noise by behavioral mode, but practical imbalance, limited capacity, and few inference steps can blur the regions.Blurred partitions may cause mode preference or collapse.Source: §1 · p. 2
- MOSS encodes semantic mode identity in Gaussian sources and adapts their separation to local multimodality, merging similar behaviors and separating them at branches.A modality-strength estimator sets the source separation while one flow network is shared across modes.Source: §1 · p. 2, Figure 1 · p. 2
- MOSS can adjust deployment mode probabilities through source sampling and share trajectory data across modes to support data-scarce behaviors.Its training-established source regions may also reduce guidance-related trajectory degradation and suit emerging one-step methods without guidance.Source: §1 · pp. 2–3
- On 3 real-world tasks, MOSS averages 84.2% GR/85.3% SR, beating mode-conditioned FM by 43.3/12.0 points; across 6 simulated tasks, steering reaches 100%.Under a 10:90 mode imbalance, minority-mode guidance rate exceeds 90%.Source: §1 · p. 3
2 Related work
Prior work addresses multimodal generation, source design, and explicit behavioral conditioning. MOSS distinguishes itself by aligning source components with semantic modes and adapting their separation to local multimodality.
Multimodal generative policies
- Multimodal generative policies can entangle or collapse behavioral modes; VFP and MARS add latent structure or adaptive generation but do not actively control multimodality.VFP combines a variational latent prior, optimal transport, and a mixture-of-experts decoder; MARS switches between stochastic and deterministic generation.Source: §2 · p. 3
Source shaping for flow matching
- Source-shaping methods target efficient or aligned transport, while SL-FM separates branches; MOSS instead semantically aligns source components and adapts separation to local multimodality.The passage describes A2A, VITA, and LAFM as targeting efficiency or transport alignment.Source: §2 · p. 3
Explicit conditioning for behavioral control
- Explicit conditioning injects behavioral intent, but can encourage shortcut reliance; MOSS controls behavior through the source distribution without exposing the mode variable to the policy.The passage reports that experiments show stronger controllability than directly injecting the mode variable into the velocity-field condition.Source: §2 · p. 3
3 Multimodal Property of Flow Matching
Deterministic flow transport gives behavioral modes distinct structure in the source space, with ideal separation but practical blurring. This structure motivates selecting behaviors through the source rather than relying on random sampling.
3 Multimodal Property of Flow Matching
- In trained FM policies, initial noises leading to different behavioral modes tend to occupy distinct but entangled source-space regions.Standard FM transports Gaussian source noise toward an observation-conditioned action distribution.Source: §3 · pp. 3–4, Figure 2 · p. 4
3.1 Noise Zoning and Its Geometric Regularity
- Under exact transport, mode-specific source zones are pairwise disjoint up to a null boundary and carry the probability mass of their corresponding behavioral modes.The zones arise as preimages of semantic action regions under the flow map.Source: §3.1 · p. 4
- Modes farther from the semantic boundary have a larger same-mode neighborhood in the source space.The result assumes a bi-Lipschitz flow map and relates source-space radius to endpoint semantic margin.Source: §3.1 · pp. 4–5
3.2 Practical FM Blurs the Noise Zoning
- Finite demonstrations, mode imbalance, limited model capacity, and finite-step integration can widen ambiguous source regions and cause mode confusion.Mode confusion is likely where transport error becomes comparable to semantic margin.Source: §3.2 · p. 5
4 Mode-Oriented Source Shaping
MOSS selects a behavioral mode by shaping its initial source distribution while retaining a shared flow field. It estimates local multimodality to adjust how separated the mode-specific sources should be.
4.2 Modality Strength Estimator
- The Modality Strength Estimator predicts high strength near behavioral divergence and lower strength after a mode is resolved.It is trained from local comparisons of between-mode and within-mode future-action variation, then calibrated using training-set percentiles and temporally persisted during rollout.Source: §4.2 · pp. 5–6
4.3 Adaptive Mode-Oriented Source Shaping
- MOSS uses divergent future-action means to define a mode-discriminative subspace for shaping the source distributions.The subspace is spanned by the mode-specific components after removing what is shared across modes.Source: §4.3 · p. 6
- As modality strength increases, mode-specific sources become more separated and less overlapping; at low strength, they stay closer to the shared Gaussian prior.The modality strength controls both source-center displacement and overlap.Source: §4.3 · pp. 6–7
4.1 Overview
- MOSS encodes the requested mode in a mode-specific source distribution, selecting behavior by changing the initial condition while keeping one shared FM velocity field.The mode label is not supplied to the velocity field.Source: §4.1 · p. 5, §4.3 · p. 7
5 Evaluation
MOSS is evaluated on simulated and real-world manipulation tasks for mode control, minority-mode learning under imbalance, and design choices. Results report strong mode controllability, including across multiple modes and real-robot tasks.
5.1 Experimental Setup
- The evaluation covers four bimodal and two multi-mode simulation tasks, plus Place Cups, Pick Veg, and Open Drawer on a Galaxea R1 Lite robot.The evaluation targets mode steering, minority-mode recovery under imbalance, and potential benefits for success and training efficiency.Source: §5.1 · pp. 7–8
5.2 Controllable Multimodal Behavior
- MOSS achieves the strongest average mode controllability across four bimodal simulation tasks while maintaining high task success; mode-conditioned FM is competitive and GMM-Sourced FM varies in semantic alignment.Table 1 reports guidance rate (GR) and success rate (SR) for each requested mode.Source: §5.2 · p. 8, Table 1 · p. 8
- MOSS maintains high success and guidance rates across multiple requested modes on 2D Navigation and Place Cube.Estimated modality strength is low when behaviors share similar motion and peaks near trajectory divergence, where MOSS increases source separation.Source: Figure 4 · p. 7, §5.2 · p. 8, Figure 6 · p. 8
- MOSS achieves strong success and guidance rates on all three real-world tasks and is more balanced across modes than mode-conditioned FM.Mode-conditioned FM fails to control mode 1 on Pick Veg and mode 2 on Open Drawer; increasing classifier-free guidance does not consistently recover both SR and GR.Source: §5.2 · p. 8, Table 2 · p. 9
5.3 Robustness and Training Efficiency
- On 10:90 imbalanced datasets, MOSS reaches high minority-mode guidance with relatively few target-mode exposures and approaches final performance earlier on Push Cube than both baselines.MOSS and mode-conditioned FM use all 100 demonstrations, while FM-DiT uses only the 10 minority-mode demonstrations; the former methods therefore also see majority-mode examples.Source: §5.3 · p. 9, Figure 7 · p. 9
5.4 Ablation Studies
- Ablations show that a shared Gaussian gives poor controllability, single-step separation helps only modestly, and full-chunk separation substantially increases guidance rate.Adaptive shaping further improves over fixed full-chunk separation and typically reaches strong performance earlier in training.Source: §5.4 · p. 9, Table 3 · p. 9
6 Conclusion
The paper presents MOSS as a source-space approach to controllable multimodal flow-matching policies. Its results support strong mode control and task success across simulated and real-world manipulation, including minority-mode learning under imbalanced demonstrations.
6 Conclusion
- MOSS explicitly associates semantic behavioral modes with source distributions and adaptively adjusts their separation according to local multimodality, without conditioning the shared velocity field on the requested mode.The paper reports strong mode controllability and task success across simulated and real-world tasks, with improved minority-mode learning under highly imbalanced demonstrations.Source: §6 · pp. 9–10
7 Limitations
MOSS currently assumes a fixed, discrete set of behavioral modes specified before training, and its performance in larger-scale policy settings has not been evaluated.
7 Limitations
- The formulation does not directly support open-ended mode control or continuous behavioral variation, and performance with larger-scale generative robot policies remains unevaluated.Dynamic mode discovery and continuous mode spaces are identified as future directions.Source: §7 · p. 10
A Visualization of Initial-Noise Mode Separation
The section examines whether initial noise in a standard flow-matching policy reflects the behavioral mode eventually executed. It reports structured but overlapping mode regions, and tests Gaussian sources centered on their mode-specific means.
A Visualization of Initial-Noise Mode Separation
- In 1,000 rollouts on four-route 2D Navigation, 939 successful trajectories were labeled by their mean vertical coordinate in the wall-crossing region.The study records the first executed action chunk's eight two-dimensional actions and uses the y-dimension as an eight-dimensional noise feature.Source: A Visualization of Initial-Noise Mode… · p. 13
- The visualization projects standardized eight-dimensional noise features into two dimensions with LDA; points are colored by resulting mode and ellipses show nominal 95% Gaussian covariance contours.LDA is fitted using trajectory-derived mode labels.Source: A Visualization of Initial-Noise Mode… · p. 13
- Mean y-noise values for Modes 1–4 increase from −0.303 to 0.286, consistent with alignment between source-space organization and route geometry.The naturally emerging mode regions substantially overlap, so this organization is insufficient for reliable control.Source: A Visualization of Initial-Noise Mode… · p. 13
- Gaussian sources centered at the Mode 1 and Mode 4 initial-noise means generate their corresponding modes in 62% and 56% of rollouts, while intermediate means yield more overlapping mode distributions.The experiment retains the original source standard deviation, σ = 1, and performs 100 rollouts per source.Source: A Visualization of Initial-Noise Mode… · p. 13, Table 4 · p. 14
B.1 Proof of Theorem 1
- The proof argues that pairwise-disjoint semantic regions have pairwise-disjoint source-space noise zones under the stated transport mapping.The passage states this disjointness result and concludes the proof of Theorem 1.Source: §B.1 · p. 14
B.2 Flow Regularity and Proof of Theorem 2
The proof establishes mode-dependent structure in source space under regularity assumptions, then bounds how endpoint error affects mode preservation. When σmin = 0, the guarantee is local rather than global.
B.2 Flow Regularity and Proof of Theorem 2
- Under global well-posedness and bounded posterior covariance, the flow has global spatial regularity; Theorem 2 then establishes connected same-mode neighborhoods and separation across modes.The proof uses a bi-Lipschitz flow property to support these geometric claims.Source: §B.2 · pp. 14–15
B.3 Proof of Proposition 3 and Corollary 4
- As the endpoint-error bound increases, the set of source points without guaranteed mode preservation expands monotonically.Outside the boundary-sensitive region, Proposition 3 guarantees preservation of the behavioral mode.Source: §B.3 · p. 16
B.4 Local Guarantee for σmin = 0
- When σmin = 0, the global sufficient bound is unavailable, but local same-mode neighborhoods and cross-mode separation hold where the velocity Jacobian is integrable along relevant trajectories.The local argument applies on compact source regions; it does not assume a global bi-Lipschitz bound.Source: §B.4 · pp. 16–17
C.1 Architecture and Hyperparameters
The method estimates local multimodality, calibrates that signal, and uses it to adapt mode-specific source geometry during flow-matching training and deployment. The supplied sections also describe the evaluation tasks, baselines, and metrics.
C.1 Architecture and Hyperparameters
- The modality-strength estimator is a two-layer MLP that uses eight-step observation and executed-action histories as input.It has hidden width 256, GELU activations, dropout 0.1, and a scalar sigmoid output.Source: §C.1 · p. 17
C.2 Neighborhood Construction and Supervision
- Local multimodality supervision uses retrieved neighbors and their weighted future-action chunks to compare between-mode and within-mode variation.The confidence weight suppresses unreliable targets from distant or poorly represented neighborhoods.Source: §C.2 · pp. 17–18
C.3 Calibration and Temporal Persistence
- Training-split calibration maps active estimator predictions to a normalized strength, while rollout persistence prevents detected multimodality from immediately losing source separation.The 95th percentile of active training predictions maps to q = 1; the persistence rule uses qfloor = 0.05 and λq = 0.75.Source: §C.3 · p. 18
D.1 Branch-Critical Region Selection.
- The branch-critical region is selected from training windows with the highest multimodality, and its mode means define the source-space mode-difference basis.These statistics are computed from training demonstrations and fixed before flow-matching training and deployment.Source: §D.1 · pp. 18–19
D.2 Adaptive Separation Schedule
- Increasing modality strength moves mode-specific source centers farther apart and increases their separation relative to spread, with clipping limiting concentration or diffusion.The orthogonal source scale remains fixed.Source: §D.2 · p. 19
D.3 Training and Inference Procedure
- The estimator and calibration are frozen before flow-matching training; training selects sources by demonstration mode, while deployment selects them by requested mode.The method uses one shared velocity field across modes.Source: §D.3 · p. 19
E.1 Tasks and Mode Definitions
- Evaluation covers six simulation tasks and three real-world tasks, with mode annotations representing distinct successful strategies rather than task success itself.The tasks include four bimodal and two multi-mode simulation tasks, plus Place Cups, Pick Veg, and Open Drawer on a real robot.Source: §E.1 · p. 19
E.5 Evaluation Protocol and Metrics
- The real-robot protocol reports guidance rate only for successful rollouts that realize the requested mode, alongside success rate.Mode-controllable policies receive 20 trials per requested mode; the real-world setup uses a Galaxea R1 Lite robot.Source: §E.4 · p. 20, §E.5 · p. 21
F Classifier-Free Guidance for Mode-Conditioned FM
The section tests whether classifier-free guidance improves mode-conditioned FM and describes the imbalanced-data and source-distribution ablation protocols. Increasing guidance helps on some real-robot modes, but the reported changes are not consistent across tasks.
F Classifier-Free Guidance for Mode-Conditioned FM
- On Place Cups, increasing CFG raises mode-1 guidance from 0.15 to 0.65; on Pick Veg, GR rises only to 0.20 while SR falls from 0.60 to 0.45.On Open Drawer, mode-2 GR rises from 0.00 to 0.25 while SR falls from 0.85 to 0.75; the authors conclude CFG cannot reliably resolve asymmetric mode collapse.Source: F Classifier-Free Guidance for… · p. 22, Table 9 · p. 22
G Imbalanced Training and Data Efficiency
- The imbalanced-data protocol trains MOSS and mode-conditioned FM with 10 mode-1 and 90 mode-2 demonstrations, then evaluates the requested minority mode.The checkpoint curves compare methods by cumulative sampled minority-mode windows; MOSS can also use majority-mode demonstrations.Source: G Imbalanced Training and Data… · p. 22
H.1 Ablation Variants
- The ablations compare a sign-partitioned single Gaussian, separation only at the branching point, fixed full-chunk separation, and modality-strength-adaptive full-chunk separation.Single-position variants leave most source dimensions shared; full-chunk variants maintain mode identity across the predicted action sequence.Source: §H.1 · p. 22
H.2 Fixed versus Adaptive Training Curves
Adaptive MOSS maintains higher guidance rates than fixed full-chunk shaping across the reported checkpoints. Offline analysis also finds reduced source alignment error and transport, especially along the mode-discriminative direction.
H.3 Source Alignment and Transport Cost
- Across all three tasks, adaptive shaping reduces Dµ and total squared transport, with the clearest D∥ differences on Push Cube and Stack Cube.The orthogonal transport component is largely shared by the two variants.Source: Table 10 · p. 23, §H.3 · pp. 23–24
H.2 Fixed versus Adaptive Training Curves
- At epoch 100, adaptive MOSS mean GR is 0.96, 0.99, and 0.98 on Push Cube, Grasp Eyeglasses, and Stack Cube, versus 0.90, 0.88, and 0.85 for fixed shaping.The checkpoint-wise comparison evaluates both requested modes every 10 epochs using 50 rollouts per mode; the adaptive gains persist across training.Source: §H.2 · p. 23