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

Abstract

Flow-matching-based vision-language-action (VLA) models have emerged as powerful policies for robotic manipulation, yet a critical capability remains underexplored: fine-grained behavioral control, the ability to govern how a robot performs a task by intervening on its internal representations. Representation steering is a well-established interpretability tool for language and vision-language models, where behavioral features are typically encoded as linear directions, but we show that these classic methods fall short in flow-matching based VLAs. We propose DiMaS, a Distribution Matching Steering strategy tailored to flow-matching VLAs, which transports between representation distributions rather than shifting along a fixed direction, and show that it effectively controls behavior across two state-of-the-art VLAs.

Introduction & Setup

Vision-Language-Action Models (VLAs): For readers familiar LLM/VLM architectures but unaware of VLAs, you can think of current SoTA VLAs architecturally as a mixture of MLLMs and flow-matching networks. At any time instance, a VLA receives as input consisting of camera images, a text instruction for the underlying task and its current joint state information. This input is processed by the VLM part of the VLA, which conditions a flow-matching based action expert that predicts the low-level actions for each robot joint.

Our goal proposes a method that can offer a user, post-hoc behavioral control over robotic actions predicted by vision-language action models (VLAs). Our method, titled DiMaS (Distribution Matching for Steering VLAs) is a representation steering method. Readers acquainted with interpretability literature might already be familiar with the term, but the essential idea is to modify the internal representations of a generative model during inference time in order to achieve desired control over its output.

Our contributions: We make key advancements along three axes

Steering with task success: The holy grail for VLA steering we believe is about how effectively a steering method can control robot style without losing on the task success rate. Even though our method is far from perfect, it takes some crucial steps forward towards achieving this goal.

Generalizability evaluation: Evaluation in VLA papers often suffer from poor generalization evaluation. In contrast, we perform a structured evaluation for how well our steering method generalizes under various splits of training/testing tasks.

VLA Representation strucutre insights: Steering in LLMs/VLMs is popularly performed by shifting in a single direction in the latent space. While the evidence against representation of a feature as a linear direction is mounting, prior VLA steering methods still followed this paradigm. We qualitatively show that it isn't really the case for common VLA architectures (despite high linear separability!) and consequently propose an optimal-transport based method.

Method overview

DiMaS matches action expert activations between two distributions (e.g. slow vs. fast, low vs. high), then steers a VLA policy toward one of the clusters at inference time

Our method details can be summarized as the following steps:

Representation and feature sourcing: We first extract representations from a chosen layer of the flow-matching action expert (not the VLM!), across a number of tasks and episodes. Each representation is paired with a target scalar feature that denotes the behavioral property we wish to control.

Source/target distribution construction: The complete set of extracted representations are partitioned into a source and target distribution based on the feature value. In the DiMaS illustration, this corresponds to the red/green distributions.

Learning the steering as transport map: The underlying foundation of our steering intervention relies on a distribution matching algorithm to learn a correspondence between the two representation source/target distributions. The main reason we opted for such a design was to ensure our steering operation could indeed move from one distribution to the other. Even though we observed a high linear separability in the action expert representations (specially in the deeper layers, later flow-matching steps), single fixed direction in the latent space was not a "good" representation of the underlying feature. Intuitively, even though the red/green points in the DiMaS illustration above are linearly separable, moving the red points by a constant shift in any direction cannot match them completely to the green points (because the underlying geometries are different). Further details about this qualitative study can be found in the paper.
Test-time intervention strategy: At inference time, we check whether the current representation already lies in the desired distribution (via a learned classifier between source/target distributions). If not, we transport it using the map learned offline. Rather than applying the full transport, we interpolate between the original layer output and the transported representation, giving smooth, controllable steering strength. This interpolation is essential to ensure we are able to control the feature of interest without significantly dropping the task success rate.

Steering Demonstrations

Baseline vs. DiMaS-steered rollouts on the same task and seed.

Speed steering, high → low

Baseline (unsteered)

Steered (high → low)

The end-effector moves noticeably slower after steering.

Height steering, low → high

Baseline (unsteered)

Steered (low → high)

The end-effector moves higher (larger |Δz| per step) after steering.

BibTeX

@misc{dimas2026,
  title  = {{DiMaS}: Distribution Matching for Steering Vision-Language-Action Models},
  author = {Khayatan, Pegah and Meziane, Sara and Parekh, Jayneel and Cord, Matthieu},
  year   = {2026},
  note   = {ECCVW 2026 / https://arxiv.org/abs/2607.14280},
}