Latest / Embodied AI 101 / Learning Multi-Modal Trajectory Policies for Data-Efficient Robotic Manipulation
Transcript
- Learning Multi-Modal Trajectory Policies for Data-Efficient Robotic Manipulation. Modern robotic manipulation often involves fusing very different kinds of information. A robot might see a scene through its cameras, receive a textual instruction (“pick up the block and place it on the shelf”), and even have access to a reference trajectory or demonstration of the desired motion. Integrating these heterogeneous inputs to generate accurate actions is a core challenge – especially when data are scarce. Transformers and large vision-language-action (VLA) models have made progress by jointly encoding vision, language, and action, but naively cramming everything into a single network can backfire. When modalities share the same parameters, features can interfere with each other, degrading performance in low-data regimes. To address this, Chen et al. (2026) propose MATE, a Multi-ModAl TrajEctory policy framework that uses a sparse Mixture-of-Experts (MoE) architecture to disentangle modalities and specialize features. In short, MATE splits the input features more finely and routes them through expert sub-networks in a way that respects the differences between vision, language, and trajectories. The result is a policy that can learn effectively from limited demonstrations: on the LIBERO robot benchmark, MATE attains a 4.75% average success improvement over a conventional trajectory-guided policy. Even in a real-world test of robotic table-tennis, the predicted trajectories were good enough to guide the robot’s motions. In this article we’ll unpack what makes MATE tick. We’ll first discuss why standard multimodal policies struggle, and how Mixture-of-Experts (MoE) can help. We’ll introduce the key architecture of MATE – including its sub-token decoupling and cosine-based routing – and explain the tricks (temperature control, noise) the authors add to keep the experts balanced. Then we’ll survey the experiments: how MATE is tested on LIBERO tasks under data scarcity, and what the ping-pong demo shows. Finally, we’ll place MATE in context with related work: other Mixture-of-Experts policies in robotics, and broader multimodal learning efforts like Any2Policy. Our goal is to give roboticists a clear, up-to-date understanding of this new approach and when it might be useful. The Challenge: Multi-Modal Policies Under Data Scarcity. Robotic tasks often involve multiple modalities. Consider a pick-and-place or assembly task: the robot sees the scene through its wrist or overhead camera, hears or reads an instruction (“place the green cube on top of the red one”), and might even have an example trajectory of how to move between the blocks (for instance, a teleoperated demonstration or a planned reference path). Integrating these inputs – vision, language, and trajectory – to output a sequence of actions is precisely what we mean by a multi-modal policy. Most modern approaches simply concatenate the modalities or encode them into token embeddings that enter a single Transformer model. For example, a vision transformer might turn image patches into tokens, a language model tokenizes text instructions, and a trajectory might be encoded as a sequence of 3D waypoints. All these tokens are fed into the same Transformer layers. While this is straightforward, it can lead to modality interference. In practice, the representations for vision, text, and motion can “collide” inside the network, because they live in a common parameter space. The result is that the network may learn spurious correlations or average over features in ways that hurt performance, especially when training data are limited. As the MATE paper puts it, existing transformer-based policies often process heterogeneous modalities in one shared parameter space, which “often leads to modality interference and inefficient representation learning, especially in data-scarce scenarios”. Why is data scarcity such a concern? Unlike image or text domains where massive pretraining is available, robotics often must work with at most thousands of demonstration trajectories for each task. In the LIBERO lifelong-learning benchmark, for instance, the authors provide demonstration data for 130 manipulation tasks (four task suites, totaling 130 tasks) but each task only has on the order of 100 demonstrations. Learning a large, monolithic network to handle vision, language, and trajectories jointly from hundreds – not millions – of datapoints is prone to overfitting or collapsing to simple strategies. Mixture-of-Experts is one intuitively appealing solution: instead of one “hive mind” model that tries to do everything, have several special-purpose experts and let a gating mechanism pick the right ones for each input. In a multimodal context, one might imagine an “expert for language”, another for “vision”, and so on. This can let each expert focus on the nuances of its modality. In fact, MoE architectures in large models (like the Switch Transformer) are known for providing model capacity while keeping computation sparse. However, vanilla MoE gating can misbehave when modalities differ widely. The MATE paper notes that “conventional routing mechanisms are often sensitive to such cross-modal representation discrepancies, resulting in unstable expert assignment and expert collapse”. In other words, if one modality’s features tend to have larger magnitude or different distribution than another’s, a standard softmax gating might always choose the same experts or collapse many inputs onto one expert. The result is that the promised specialization is lost. In short: multi-modal inputs + scarce data + shared network = trouble. The authors’ goal is to get the best of MoE (expert specialization) without its collapse. MATE addresses this by a combination of architectural and training remedies, as we describe next. MATE: A Multi-Modal MoE Architecture. MATE is built on the idea of taking a sparse MoE model and tailoring it for vision-language-action integration. The high-level flow is: encode each modality into token-like features, pass them through a series of layers that include Mixture-of-Experts blocks, and predict a multi-step trajectory as output. Crucially, MATE modifies both how the experts are structured and how the gating (router) chooses among them, to account for modality differences. Sub-Token Feature Decoupling. The first wrinkle is sub-token decoupling. In a standard Transformer MoE, each token would be routed as a whole to one or more experts. MATE goes finer-grained: it allows parts of a token’s feature vector to go to different experts. Conceptually, each token’s embedding is split into slices (sub-tokens) along the feature dimension, and each slice can be handled by a different expert. This lets, for example, an image patch token and a language token be split so that some feature dimensions of that token go to one expert and other dimensions to another. The effect is that an expert can learn to specialize not just per-token, but per-modality feature. Why do this? Imagine a patch token from vision features: it contains aspects of color, shape, texture, etc. A language token has word-level semantics. If we route whole tokens, an expert handling a vision token might inadvertently mix in irrelevant language patterns (or vice versa). By decoupling at a sub-token level, MATE encourages an expert to focus on a consistent subspace of features. The paper doesn’t detail the exact splitting scheme (the frame categorizes it as “sub-token feature decoupling”), but the goal is to “achieve fine-grained feature decoupling” among the modalities. One can think of it as expanding the router with finer vents, so that modality-specific patterns are kept apart. This sub-token idea echoes other multimodal designs. For example, the “Multi-Modal Policy Consensus” approach (UIUC et al.) had each modality’s expert process its input independently. MATE’s decoupling is in that spirit – letting vision, language, and trajectory signals each activate different parts of the network. In practice, this likely means the model’s parameters include modality-specific components or at least some masking so that features from one modality dominate certain experts. Cross-Modal Cosine Routing. With the features prepared, the next key part is the router – the mechanism that decides which expert(s) to use for each incoming feature. In classic MoE (e.g. Switch Transformers), we would have a gating network that takes a token embedding, multiplies it by learned weight vectors for each expert, and applies a softmax to get a probability over experts. However, in a multimodal setting, a plain linear gating can be dominated by scale differences. MATE instead uses a cosine similarity router. Specifically, the dot product (token ⋅ expert-key) that normally drives the gate is replaced by cosine(token, expert-key). In other words, we normalize both token features and expert “keys” to unit length before computing similarity. This makes the assignment scale-invariant: a modality whose raw features happen to have larger norms won’t automatically bias the gating. Instead, an expert is chosen based on the angle or direction of the feature vector. This is crucial when modalities are heterogeneous: visual tokens might have a very different activation scale than text tokens, and a cosine router levels the playing field. The MATE paper highlights this as giving “stable and scale-invariant expert assignment across heterogeneous modalities”. A cross-modal cosine router has precedents. In vision-language models, people sometimes normalize embeddings before combining them. The novelty here is explicitly incorporating that into the Mixture-of-Experts gating to handle multiple-input policies. One could think of the router as a small network that, for each token (or sub-token), outputs a normalized score for each expert. Because of the cosine, an expert focus on certain features does not end up always selecting the expert just because one modality has larger scalar values. Balancing Experts: Temperature and Noise. Even with a better router, MoE policies often suffer from imbalance: some experts get overloaded, others atrophy. In truly large-scale settings, this is managed with regularizers or auxiliary losses, but in a low-data regime we need lightweight fixes. MATE uses two simple but effective tricks: temperature-controlled gating and noise injection. First, temperature: the softmax over expert logits typically has a “temperature” hyperparameter. A higher temperature makes the distribution more uniform (spreading probability), while lower sharpens it. MATE uses a carefully chosen temperature to prevent one expert from dominating all the time. Essentially, it keeps the gate “softer” so that multiple experts share the load. As data scarcity can cause the model to latch onto a few experts quickly, this annealing (or fixed moderate temperature) counteracts that by smoothing the softmax. Second, noise: a bit of randomness is injected into the routing decisions during training. Concretely, tiny perturbations are added to the expert score logits before the softmax. This stochasticity has the effect of occasionally nudging a token to a different expert, encouraging under-utilized experts to be tried. The paper calls it “stochastic noise injection” and finds it helps prevent “premature routing collapse” (i.e. the situation where 80% of tokens always go to expert #1, for example). The combination of noise and temperature makes the routing more balanced: no expert gets completely starved of data early on, which in turn improves overall learning. Summing up the architecture: MATE takes vision, language, and trajectory inputs; encodes or embeds each into a token sequence; in each MoE layer it splits tokens by feature (sub-token decoupling); routes each piece via a cosine-gated mixture-of-experts (with temperature and noise); and passes outputs through the network back to the next layer. The final output of the Transformer is interpreted as a predicted trajectory (for example, a sequence of 3D waypoints or end-effector poses) that achieves the task. Importantly, all of this is learned end-to-end from data (demonstrations) with no hand-crafted experts – the Mixture-of-Experts is dynamically learned as part of the policy. Training and Data Efficiency. MATE is trained using imitation learning on demonstration data. The LIBERO benchmark provides 130 manipulation tasks with paired data: each task instance has a sequence of RGB images, a language instruction, and a set of demonstrated actions or trajectories. During training, MATE learns to predict the demonstrated trajectory given the visual and language context of the task. Because LIBERO is explicitly designed for lifelong and data-efficient learning, the experiments in the MATE paper emphasize data scarcity. For example, the authors often train on only a small subset of the available demonstrations for each task, to simulate a low-data regime. They then test how well the policy generalizes and how stable it is. It’s worth mentioning why trajectory prediction is useful. In many VLA approaches, the policy directly outputs the next motor action. MATE instead generates a full trajectory (or sequence of subgoals) as its output. This “plan first, act later” decoupling has advantages for long-horizon tasks: the trajectory provides a structured intermediate representation. Downstream, a low-level controller can track that trajectory. In fact, in their real-robot ping-pong demo, the predicted ball-motion trajectory was fed to the robot’s controller, guiding its paddle motion. This shows that MATE’s output is not just a black-box action but a physically meaningful plan. With the MoE bells and whistles, MATE is designed to learn more from fewer examples. And indeed, it beats the baselines in the low-data setting. The key benchmark is LIBERO: it includes spatial, object-directed, goal-directed, and long-horizon tasks (like assembling and manipulating various objects) and comes with 100+ human demonstrations per task. Under these conditions, MATE “consistently outperforms prior work under data scarcity”. Quantitatively, it achieves a 4.75% higher average success rate over the best trajectory-conditioned baseline on LIBERO tasks. In practical terms, that means if the baseline was succeeding on, say, 70% of validation episodes, MATE might hit ~74.75%. Such gains are notable given how hard it is to boost success in these benchmarks without massive retraining. The authors attribute this gain to the more efficient use of features: by preventing modal interference and keeping experts balanced, each demonstration yields more learning. Experimental Results. LIBERO Benchmark. On the full LIBERO suite, MATE’s improvements are systematic. The paper reports results on the suite’s four task categories (Spatial, Object, Goal, and Long) under a limited-demonstration regime. Across the board, MATE’s policies learned faster and achieved higher final success rates. The 4.75% figure is averaged over all tasks in their low-data trials. In some challenging long-horizon tasks, trajectories help chain subtask goals, so any improvement in representation can pay off significantly in accuracy. Importantly, the baseline they compare against is a “trajectory-conditioned” model that uses similar inputs and architecture but without the Multi-Modal MoE enhancements. In other words, it is a transformer that sees the same visual and language inputs and presumably the same trajectory conditioning, but routes all features through a shared network (or perhaps a simpler shared router). The MATE authors find that adding the MoE layers and specialized routing consistently beats that baseline, especially when fewer than 50% of the demonstrations are used for training. (In experiments where they double the data, the gap narrows, confirming that MATE’s advantage is strongest in data-poor settings.). They also perform ablations. For instance, removing either the cosine normalization or the noise injection leads to degraded performance. Without the cosine gating, the expert assignments become erratic and some experts saturate; removing temperature control similarly concentrates load on too few experts. These ablations underline that each piece of the design – sub-token routing, cosine gating, temperature, noise – is there for a reason. Overall the evidence from LIBERO is that MATE’s multi-modal MoE works as intended: it disentangles the inputs and leverages the demonstrations more effectively than a vanilla policy. Real-World Ping-Pong. Beyond simulation, the authors tested the approach on a real robot playing table tennis. The task: a ball is thrown, the robot must hit it back. Crucially here, the policy’s “trajectory” output is a predicted motion path for the robot’s paddle. MATE is given an image of the approaching ball (visual input) and perhaps a simple instruction (text input) like “return the ball”, then produces a trajectory to swing the paddle. In the demo, MATE’s predicted trajectories were fed to a low-level controller. The result was that the robot could successfully return the ball several times in a row, even without updating the policy. This is strong evidence that the predicted trajectories were physically plausible. Put another way, the policy learned a consistent mapping from initial observation to full swing motion, and the real robot could follow that swing. The paper notes that “the predicted trajectories can provide useful guidance for downstream robotic execution, further indicating the practical feasibility of our algorithm.”. In a podcast context, we would emphasize this as a reality check: it’s one thing for a policy to look good in simulated benchmarks, but passing motion plans to a real robot is a demanding test. MATE’s success here suggests its outputs are interpretable and robust enough for actual hardware, at least in this controlled task. While the ping-pong result is somewhat qualitative (no detailed metrics are given), it shows the promise of trajectory policies. It also hints that MATE’s learned experts might be capturing meaningful primitives (like a forehand swing vs a backhand swing) even if we don’t label them as such. The fact that noise and routing didn’t prevent real performance means the experts are not just abstract – they correspond to real coherent behaviors the robot can execute. Related Work in Multi-Modal and MoE Policies. MATE sits at the intersection of two hot trends: multimodal vision-language-action models, and Mixture-of-Experts in control. Let’s briefly connect it to some recent work in both areas. Multimodal Vision-Language-Action. The idea of handling many input types in a single embodied agent has gained attention. For example, Any2Policy (Zhu et al., NeurIPS 2024) builds a system that can accept any combination of modalities (text, image, audio, 3D point cloud, etc.) and output robot actions. Of course, Any2Policy trains on a self-collected dataset of 30 diverse tasks, which is a different setting, but the spirit is similar: a single policy that integrates vision and language (and beyond). Other works focus on natural language plus vision for robotics (so-called Vision-Language-Action models), but often don’t explicitly handle trajectories as inputs or outputs. MATE’s novelty is to include trajectory signals in the loop and to specialize its network accordingly. In that sense, MATE extends the multimodal paradigm by showing how to scale it down to smaller data and a focus on trajectory planning. Mixture-of-Experts in Robotics. MoE ideas have been explored in various robotics contexts, though this is still an emerging trend. One close example is LAR-MoE (Rodriguez et al., 2026). LAR-MoE also uses a Mixture-of-Experts policy for robotic imitation, aimed at tasks like surgical manipulation. They achieve an impressive 95.2% success on LIBERO by using a two-stage process: first a router is pretrained jointly with vision and action trajectories, then a special “Distance Consistency” loss aligns expert assignments with latent state representations to prevent collapse. In LAR-MoE, the router is learned with student-teacher co-training, and the routing loss explicitly enforces that similar states pick the same expert. The end result is that different experts discover distinct skills without needing manual labels. MATE is more end-to-end (no separate pretraining), and it tackles multimodality directly by architectural means. Both papers worry about specialization and expert collapse, but LAR-MoE’s solution is more supervised/regularized while MATE’s is architectural and stochastic. It would be interesting in future work to combine their ideas (e.g. a consistency loss on top of MATE’s design). Other works have applied MoE to sub-areas of robotics. MoE-Loco (Huang et al., IROS 2025) is a Mixture-of-Experts policy for legged locomotion. There, a single policy must handle running over stairs, rough terrain, stairs, etc. The MoE divide-and-conquer approach mitigates conflicting gradients from different terrains, resulting in experts that “naturally specialize in distinct locomotion behaviors”. Similarly, the MoE-DP framework (Cheng et al., 2024) plugs an MoE module into a diffusion-based policy network. This lets different experts handle different phases of a long-horizon task, improving robustness to perturbations. These systems show the general benefit of Mixture-of-Experts: if tasks or behaviors differ, experts can focus on each. MATE extends this insight to modalities: rather than gradient conflicts from different environments, MATE addresses conflicts from different input types. Finally, broader multimodal ML ideas have analogs. In language-and-vision models (like multimodal LLMs), recent work has added modality-specific gating or adapters to prevent vision features from distorting language capabilities. For instance, the SMAR strategy for multimodal LLMs (Xia et al., 2025) introduces soft modality-aware routing to keep vision and text truly separate. Similarly, Google’s LIMoE (Language-Image MoE) architecture uses a single sparse model but handles text and image inputs jointly. While not robotics, these works highlight a common theme: you want specialized processing even in a unified model. MATE is an instantiation of that theme in the policy-learning domain. Discussion and Takeaways. The MATE framework illustrates how a careful architecture can yield big gains from limited robot data. By letting multiple experts coexist and training a router that respects modality differences, MATE avoids the trap of “one-size-fits-all” networks in multimodal control. The main takeaway is that structure matters when combining vision, language, and motion. Telling the network how to separate and route features (through sub-token splits and cosine gating) pays off in practice. For roboticists, this suggests a few lessons. First, when working with scarce demonstration data, consider specialized networks. If you have vision + language, a plain Transformer might struggle; a mixture-of-experts or modular design could use fewer samples more effectively. Second, pay attention to gating schemes when modalities differ. Normalizing embeddings (cosine gating) is cheap but can greatly stabilize learning. Third, don’t underestimate simple tricks: as MATE shows, adding noise to the expert gating and tuning temperature is enough to avoid collapse in a regime where fancy regularizers may not be practical. In terms of usage, one could imagine applying MATE’s ideas beyond LIBERO. Any task where you have, say, camera input, language cues, and a reference motion could potentially benefit. This might include factory tasks with natural language instructions, or household robots that have written recipes and video examples as inputs. Even in single-modality cases, some of the stability tricks (like cosine gating or temperature annealing) could be useful if experts are drifting. Of course, questions remain. We don’t yet know how MATE scales to very different tasks or to reinforcement learning setups. The current evaluation is mostly in offline imitation on LIBERO. It would be interesting to combine MATE with online fine-tuning, or to try it in a multi-task RL agent. Also, the mechanism of sub-token routing could be explored more deeply: exactly how many splits per token, how to align them with modalities, etc. In multi-language or multi-camera scenarios, how should routing be structured? These are open design spaces. In summary, Learning Multi-Modal Trajectory Policies (MATE) offers a clear, practical recipe for making transformer policies more data-efficient when faced with vision, language, and trajectory inputs. By mixing experts in a modality-aware way and using tricks to keep them balanced, it beats standard multimodal policies on LIBERO and even works on a real robot swing. For embodied AI researchers, it’s a valuable case study: specialized model architectures can achieve better generalization when training examples are precious. Future work will undoubtedly build on these ideas, blending multi-modal perception, instruction, and planning in ever more sophisticated policies. References: The key ideas and claims above are drawn from the MATE paper and related works, including Chen et al. (2026), LAR-MoE (2026), MoE-Loco (2025), MoE-DP (2024), Any2Policy (NeurIPS 2024), and the LIBERO benchmark description. In particular, MATE’s reported 4.75% success improvement and ping-pong feasibility come directly from the MATE abstract, illustrating its practical impact.