🤖 machine learning

DiMaS: Distribution Matching for Steering Vision-Language-Action Models

This paper introduces DiMaS, a distribution-matching steering strategy that overcomes the limitations of classical linear methods to enable fine-grained behavioral control in flow-matching-based vision-language-action models by transporting representation distributions rather than shifting along fixed directions.

Original authors: Pegah Khayatan, Sara Meziane, Jayneel Parekh, Matthieu Cord

Published 2026-07-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Pegah Khayatan, Sara Meziane, Jayneel Parekh, Matthieu Cord

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where robots don't just follow a rigid script, but can actually "feel" the difference between a gentle nudge and a forceful shove. This is the realm of Vision-Language-Action (VLA) models, a cutting-edge corner of artificial intelligence where computers learn to see the world through cameras, understand instructions through language, and then physically move robot arms to get things done. Think of these models as a super-smart brain that has read every book and watched every video, now trying to figure out how to use a pair of mechanical hands.

But there's a catch: while these robots are getting better at what to do (like "pick up the cup"), they are still a bit clumsy about how to do it. We want to be able to tell them, "Pick up the cup, but do it slowly," or "Pick it up, but don't lift it too high." In the world of AI, this is called behavioral control. For a long time, scientists tried to control these robots by simply nudging their internal "thoughts" in a straight line, like pushing a car slightly to the left to make it turn. However, this paper suggests that for the newest, most advanced robots, that straight-line push isn't enough. It's like trying to steer a complex dance by just telling the dancer to "move left" when the dance actually requires a spin, a dip, and a leap all at once.

Enter DiMaS (Distribution Matching for Steering), a new method proposed by Pegah Khayatan and her team. Instead of just pushing the robot's internal thoughts in a single direction, DiMaS acts like a master choreographer who looks at the entire group of "slow" movements and the entire group of "fast" movements. It learns the exact shape of the "slow" group and the "fast" group, then gently reshapes the robot's current thoughts to fit into the "slow" group's pattern.

The researchers tested this on two of the smartest robot brains available today, named SmolVLA and 𝜋0.5, using a virtual playground called LIBERO where robots have to perform tasks like stacking blocks or moving objects. They found that the old "straight-line" methods were inconsistent; sometimes they worked, but often they failed to change the robot's speed or height, or they made the robot fail the task entirely. In contrast, DiMaS successfully taught the robots to move slower or lift objects lower, and it did so while keeping the robots on track to finish their jobs.

The team also discovered something fascinating about why the old methods failed. They peered inside the robot's brain and found that while you can easily tell the difference between a "fast" thought and a "slow" thought (they are easy to separate), you cannot turn one into the other just by adding a simple number. The relationship is too complex, like trying to turn a square into a circle just by stretching it in one direction; you need to reshape the whole thing. DiMaS does exactly that: it reshapes the robot's internal state to match the desired behavior.

Perhaps most impressively, the researchers tested if this new skill could transfer to tasks the robot had never seen before. They taught the robot to move slowly using one set of puzzles and then asked it to solve completely different puzzles. The method held up surprisingly well, suggesting that the robot learned a general rule about "moving slowly" rather than just memorizing a specific trick for one specific game. While the robot did struggle a bit more when the tasks became very different or when the robot had to move for a very long time, the ability to control how the robot moves without breaking its ability to finish the task is a significant step forward.

In short, this paper suggests that to truly control advanced robots, we need to stop treating their internal minds like simple dials we can just turn up or down. Instead, we need to treat them like complex landscapes, using a map (DiMaS) to guide them from one type of behavior to another, ensuring they can still dance the whole way through.

Technical Summary: DiMaS – Distribution Matching for Steering Vision-Language-Action Models

Problem Statement

Vision-Language-Action (VLA) models, particularly those utilizing flow-matching action experts, have emerged as powerful policies for robotic manipulation. However, a critical capability remains underexplored: fine-grained behavioral control. While representation steering is a well-established interpretability tool for Large Language Models (LLMs) and Vision-Language Models (VLMs)—where behavioral features are typically encoded as linear directions—this paper demonstrates that classic linear steering methods fall short in the VLA setting.

The core challenge is that VLAs generate continuous action trajectories via flow-matching rather than discrete tokens. The authors hypothesize that behavioral features in these visuomotor settings are not linearly steerable in the same way they are in language models. Existing methods that rely on fixed additive shifts along a linear direction fail to effectively control specific behavioral attributes (such as motion speed or vertical displacement) without compromising task success or failing to induce the desired change reliably.

Methodology: DiMaS

The authors propose DiMaS (Distribution Matching Steering), a strategy tailored specifically for flow-matching VLAs. Instead of shifting representations along a fixed linear direction, DiMaS learns a mapping that transports representations between distributions.

Core Components

Representation Sourcing: The method extracts residual stream representations ( $h$ ) from the action expert ( $f_A$ ) of the VLA at specific layers ( $l$ ) and flow-matching denoising steps ( $m$ ). The action expert is chosen over the VLM backbone because representations in the action expert exhibit higher linear separability for behavioral features.
Feature Extraction: A continuous scalar behavioral feature (e.g., end-effector speed or vertical displacement) is derived from the predicted actions.
Distribution Partitioning: Representations are partitioned into a source distribution ( $\mathcal{D}^-$ ) where the feature is absent (low values) and a target distribution ( $\mathcal{D}^+$ ) where the feature is present (high values). This is achieved by thresholding on empirical quantiles to ensure well-separated populations.
Optimal Transport Map: The steering intervention is defined as a transport map $\mathcal{T}$ that moves representations from $\mathcal{D}^-$ to $\mathcal{D}^+$ . This map is learned by minimizing the Kantorovich optimal transport objective (specifically the $W_2$ distance) using a low-rank Sinkhorn algorithm. This respects the full geometry of the distributions rather than assuming a single linear direction.
Inference-Time Intervention:
- Gating: A binary classifier (linear probe) determines if a representation belongs to the source distribution (feature absent). Steering is only applied if the feature is absent.
- Interpolation: To maintain task success, the method does not fully transport the representation to the target distribution. Instead, it interpolates between the original representation and the transported representation using a knob $\alpha \in [0, 1]$ .
- Update Rule: $h' = (1-\alpha)h + \alpha(\mathcal{T} \circ P(h))$ , where $P$ is a nearest-neighbor projection onto the source distribution.

Key Contributions

DiMaS Mechanism: A novel steering strategy based on distribution matching (optimal transport) specifically adapted for flow-matching VLAs. It intervenes on latent representations and is validated across models of different scales (SmolVLA and $\pi0.5$ ).
Generalizability Analysis: The paper characterizes the transferability of behavioral control across increasingly dissimilar tasks. It demonstrates that steering learned from a subset of tasks can generalize to held-out tasks and even disjoint task suites, though effectiveness varies with task diversity.
Structural Insights: Through an analysis of the action expert's representation structure, the authors explain why classical linear steering fails in VLAs. They show that while behavioral features are linearly decodable (separable by a classifier), they are not linearly steerable (a fixed shift cannot transform the source distribution into the target distribution). This motivates the distribution-matching design of DiMaS.

Experimental Results

The method was evaluated on the LIBERO benchmark (Object, Spatial, Goal, and 10 suites) using two state-of-the-art VLAs: SmolVLA and $\pi0.5$ .

Control Efficacy: DiMaS successfully modulated target features (motion speed and vertical displacement) in both increasing and decreasing directions. In contrast, linear steering baselines (mean-difference and regression-based) and prompt injection baselines performed inconsistently, often failing to shift the feature or shifting it in the wrong direction.
Task Success: DiMaS largely preserved task success rates when modulating speed. However, modulating vertical displacement incurred a drop in success, which the authors attribute to the feature being intrinsically tied to task completion (e.g., insufficient vertical lift prevents object placement).
Generalization:
- Task Diversity: Learning from multiple tasks jointly (Setting 2) yielded more stable effects than learning from a single task (Setting 1).
- Out-of-Distribution Transfer: DiMaS remained effective when applied to held-out tasks within a suite and, to a lesser extent, when transferred to disjoint suites. The effectiveness was weakest in the most diverse suite (LIBERO-Goal), suggesting that high behavioral heterogeneity challenges the learned transport map.
- Long-Horizon Tasks: The method remained effective on longer-horizon tasks (LIBERO-10), showing a significant reduction in speed with a manageable cost to success rate.
Interpolation Trade-off: The interpolation parameter $\alpha$ allows for a tunable trade-off between steering strength and success rate. An intermediate value (e.g., $\alpha=0.5$ ) was found to offer a favorable balance, substantially reducing speed while preserving most successful episodes.

Significance and Claims

The paper claims that DiMaS provides a robust mechanism for post-hoc behavioral control in VLAs, addressing a gap where users can shape how a robot performs a task, not just what it does.

Theoretical Insight: The work provides a crucial distinction in the VLA domain: behavioral features are linearly decodable but not linearly steerable. This finding challenges the direct application of linear steering paradigms from LLMs to visuomotor control and justifies the use of optimal transport.
Practical Utility: By preserving task success while controlling specific behaviors (like speed), DiMaS offers a practical interface for shaping robot behavior without retraining the model.
Limitations and Modesty: The authors acknowledge that steering vertical displacement is inherently risky for task success due to physical constraints. They also note that the current method treats all timesteps uniformly and that the association between representations and actions in chunked outputs is an approximation.

In summary, DiMaS demonstrates that distribution matching is a necessary and effective approach for steering flow-matching VLAs, offering a generalizable framework for representation-level control in robotic manipulation.