深入理解SmolVLA中的Flow Matching Action Expert：从理论到实现-EW帮帮网

本文详细介绍了SmolVLA中Flow Matching Action Expert的理论基础、架构设计和代码实现。希望能够帮助读者深入理解这一创新技术，并为相关研究和应用提供参考。

引言

在机器人学习领域，如何让机器人理解自然语言指令并执行相应的动作一直是一个核心挑战。最近，视觉-语言-动作(Vision-Language-Action, VLA)模型的出现为这一问题提供了新的解决思路。其中，SmolVLA作为一个小型高效的VLA模型，在其架构设计中采用了一个创新的组件——Flow Matching Action Expert，这个组件负责将视觉和语言信息转化为具体的机器人动作序列。

想象一下，当你告诉机器人"把红色的苹果放到篮子里"时，机器人需要经历三个步骤：首先理解你的语言指令，然后观察环境中的物体，最后生成一系列精确的动作来完成任务。在这个过程中，Flow Matching Action Expert就像是机器人的"动作规划师"，它接收来自视觉和语言模块的信息，然后生成一系列连续的动作指令。

本文将深入探讨SmolVLA论文[1]中3.1 Model architecture部分提到的Flow matching action expert，从基础的数学原理开始，逐步解释其工作机制，并结合具体的代码实现来帮助读者全面理解这一技术。

一、Flow Matching基础理论

连续变换

在理解Flow Matching之前，让我们先从一个生活中的例子开始。想象你有一团橡皮泥，你想把它从一个圆形慢慢变成一个正方形。在这个变形过程中，橡皮泥的每一个点都在移动，而且这个移动过程是连续的、平滑的。

现在，如果我们用数学的语言来描述这个过程：

圆形是我们的"初始状态"，用数学符号表示为 $x_0$
正方形是我们的"目标状态"，用数学符号表示为 $x_1$
变形过程中的每个中间状态可以用 $x_t$ 表示，其中 $t$ 是时间参数，从0变化到1

这里的关键概念是：我们需要找到一个"变形规则"，告诉橡皮泥上的每一个点在每个时刻应该朝哪个方向移动，移动多快。这个"变形规则"在数学上就叫做向量场，用符号 $u_t(x_t)$ 表示。

常微分方程描述连续变换

这个连续变形过程可以用一个常微分方程(ODE)来精确描述：

$dxtdt=ut(xt)\frac{dx_t}{dt} = u_t(x_t)$

让我们逐个解释这个公式中的每个符号：

$x_t$ ：表示在时间 $t$ 时刻的状态（比如橡皮泥上某个点的位置）
$dxtdt\frac{dx_t}{dt}$ ：这是微积分中的导数记号，表示 $x_t$ 随时间 $t$ 的变化率，也就是"速度"
$u_t(x_t)$ ：向量场函数，输入当前时间 $t$ 和当前位置 $x_t$ ，输出在这个时空点应该移动的方向和速度

这个公式的物理含义是：在任何时刻 $t$ 和任何位置 $x_t$ ，物体的移动速度等于向量场在该点的值。

举个具体例子：如果 $u_t(x_t) = (2, 1)$ ，这意味着在当前位置，物体应该以每单位时间向右移动2个单位、向上移动1个单位的速度运动。

传统方法的计算挑战

在机器学习中，我们通常有大量的数据点需要从初始分布变换到目标分布。传统的连续归一化流(CNFs)方法需要在训练过程中反复求解上述ODE，这就像是每次都要完整地模拟整个变形过程。

具体来说，传统方法需要：

给定初始状态 $x_0$ 和向量场 $u_t$
通过数值积分求解ODE得到最终状态 $x_1$
计算 $x_1$ 与真实目标的差距
调整向量场参数，重复上述过程

这个过程计算成本极高，就好比你每次想知道橡皮泥变形的结果，都要亲手完整地捏一遍。

Flow Matching的核心创新

Flow Matching的创新之处在于提供了一种"无仿真"的训练方法。它不需要每次都完整地求解ODE，而是直接学习向量场 $u_t(x_t)$ 。

这就像是：我们不需要每次都完整地变形橡皮泥，而是直接学习"在任何给定的时刻和位置，应该朝哪个方向移动"这个规则本身。

二、条件流匹配(Conditional Flow Matching)

问题的核心挑战

直接学习向量场 $u_t(x_t)$ 面临一个根本问题：我们不知道真实的向量场应该是什么样的。这就像是我们想学习橡皮泥变形的规则，但我们不知道"正确答案"是什么。

条件流匹配(Conditional Flow Matching, CFM)巧妙地解决了这个问题。它的核心思想是：虽然我们不知道整体的向量场，但我们可以为每一对起点和终点设计一个简单的路径。

条件概率路径的数学表达

假设我们知道一个数据点的起点 $x_0$ 和终点 $x_1$ ，我们可以设计一条从 $x_0$ 到 $x_1$ 的简单路径。最直观的选择是直线路径：

$xt=(1−t)⋅x0+t⋅x1x_t = (1-t) \cdot x_0 + t \cdot x_1$

让我们详细解释这个公式：

当 $t = 0$ 时： $x0=(1−0)⋅x0+0⋅x1=x0x_0 = (1-0) \cdot x_0 + 0 \cdot x_1 = x_0$ （起点）
当 $t = 1$ 时： $x1=(1−1)⋅x0+1⋅x1=x1x_1 = (1-1) \cdot x_0 + 1 \cdot x_1 = x_1$ （终点）
当 $t = 0.5$ 时： $x0.5=0.5⋅x0+0.5⋅x1x_{0.5} = 0.5 \cdot x_0 + 0.5 \cdot x_1$ （中点）

这个公式描述了一条从起点到终点的直线路径，参数 $t$ 控制我们在这条路径上的位置。

对应的向量场（即移动方向）可以通过求导得到：

$ut(xt∣x1)=ddt[(1−t)⋅x0+t⋅x1]=−x0+x1=x1−x0u_t(x_t | x_1) = \frac{d}{dt}[(1-t) \cdot x_0 + t \cdot x_1] = -x_0 + x_1 = x_1 - x_0$

这个结果非常直观：无论在路径的哪个位置，移动方向都是从起点指向终点的向量。

边际向量场的构造

现在我们有了每一对起点-终点的条件向量场，但我们需要的是整体的向量场。CFM通过取期望来解决这个问题：

$ut(x)=Ex1∼q1[ut(x∣x1)]u_t(x) = \mathbb{E}_{x_1 \sim q_1}[u_t(x | x_1)]$

让我们逐步解释这个公式：

$q_1$ 是目标数据分布（比如我们想要生成的图像分布）
$x1∼q1x_1 \sim q_1$ 表示从目标分布中采样一个数据点
$u_t(x | x_1)$ 是给定终点 $x_1$ 时的条件向量场
$E[⋅]\mathbb{E}[\cdot]$ 表示期望值，即对所有可能的终点取平均

这个公式的含义是：在位置 $x$ 和时间 $t$ ，整体向量场等于所有可能条件向量场的加权平均，权重由目标分布决定。

高斯概率路径

在实际应用中，纯粹的直线路径可能过于简单。SmolVLA采用了高斯概率路径，在直线路径的基础上加入了随机性：

$xt=(1−t)⋅x0+t⋅x1+σt⋅ϵx_t = (1-t) \cdot x_0 + t \cdot x_1 + \sigma_t \cdot \epsilon$

新增的符号解释：

$σt\sigma_t$ ：时间相关的噪声水平，通常随时间递减
$ϵ∼N(0,I)\epsilon \sim \mathcal{N}(0, I)$ ：标准正态分布的随机噪声
$N(0,I)\mathcal{N}(0, I)$ 表示均值为0、协方差矩阵为单位矩阵的多维正态分布

这个公式的含义是：在直线路径的基础上，加入一些随机扰动。这种设计有两个好处：

增加路径的多样性，避免所有数据点都走完全相同的路径
提高模型的鲁棒性，使其能够处理带有噪声的输入

对应的条件向量场为：

$ut(xt∣x1)=x1−(1−σt2)xtσt2u_t(x_t | x_1) = \frac{x_1 - (1-\sigma_t^2) x_t}{\sigma_t^2}$

当噪声水平 $σt\sigma_t$ 很小时，这个公式近似为：

$ut(xt∣x1)≈x1−xtu_t(x_t | x_1) \approx x_1 - x_t$

这意味着向量场指向从当前位置到目标位置的方向，这是一个非常直观的结果。

Flow Matching的训练目标

有了条件向量场的概念，我们现在可以定义Flow Matching的训练目标。我们希望训练一个神经网络 $vθ(xt,t)v_\theta(x_t, t)$ 来近似真实的向量场 $u_t(x_t)$ 。

训练损失函数为：

$L(θ)=Et,x0,x1,ϵ[∥vθ(xt,t)−ut(xt∣x1)∥2]\mathcal{L}(\theta) = \mathbb{E}_{t, x_0, x_1, \epsilon}[\|v_\theta(x_t, t) - u_t(x_t | x_1)\|^2]$

让我们详细解释这个损失函数的每个组成部分：

$θ\theta$ ：神经网络的参数（权重和偏置）
$vθ(xt,t)v_\theta(x_t, t)$ ：神经网络预测的向量场，输入是当前状态 $x_t$ 和时间 $t$
$u_t(x_t | x_1)$ ：真实的条件向量场（我们的"标准答案"）
$∥⋅∥2\|\cdot\|^2$ ：L2范数的平方，即欧几里得距离的平方，用来衡量预测值和真实值之间的差距
$Et,x0,x1,ϵ[⋅]\mathbb{E}_{t, x_0, x_1, \epsilon}[\cdot]$ ：对所有随机变量取期望

期望符号下的随机变量含义：

$t$ ：从某个分布（如均匀分布或Beta分布）中采样的时间
$x_0$ ：从噪声分布（如标准正态分布）中采样的起点
$x_1$ ：从真实数据分布中采样的终点
$ϵ\epsilon$ ：用于构造高斯概率路径的随机噪声

这个损失函数的训练过程可以理解为：

随机选择一个时间点 $t$ 和一对起点-终点 $x_0, x_1)$
根据高斯概率路径计算当前状态 $x_t$
计算真实的条件向量场 $u_t(x_t | x_1)$
让神经网络预测向量场 $vθ(xt,t)v_\theta(x_t, t)$
计算预测值和真实值之间的差距，并更新网络参数

三、SmolVLA中的Flow Matching Action Expert

现在我们理解了Flow Matching的基本原理，让我们看看SmolVLA是如何将这一理论应用到机器人动作生成中的。

SmolVLA的整体架构

SmolVLA采用了一种模块化的设计，将复杂的视觉-语言-动作任务分解为两个相对独立的部分：

视觉语言模型(VLM)：负责理解图像和语言指令，提取高层次的语义特征
Flow Matching Action Expert：接收VLM的特征，生成连续的动作序列

这种设计的优势在于分工明确：VLM专注于回答"做什么"的问题，而Action Expert专注于解决"怎么做"的问题。

在代码实现中，这种架构体现在VLAFlowMatching类的设计上：

class VLAFlowMatching(nn.Module):
    def __init__(self, config):
        super().__init__()
        # VLM with Expert模型：结合视觉语言模型和动作专家
        self.vlm_with_expert = SmolVLMWithExpertModel(
            model_id=self.config.vlm_model_name,
            freeze_vision_encoder=self.config.freeze_vision_encoder,
            train_expert_only=self.config.train_expert_only,
            # ... 其他配置参数
        )
        
        # 各种投影层，用于特征空间的转换
        self.state_proj = nn.Linear(...)      # 状态投影
        self.action_in_proj = nn.Linear(...)  # 动作输入投影  
        self.action_out_proj = nn.Linear(...) # 动作输出投影

动作序列的噪声化过程

在SmolVLA中，机器人的动作序列被表示为 $At=(at,at+1,…,at+n)A_t = (a_t, a_{t+1}, \ldots, a_{t+n})$ ，其中每个 $a_i$ 是一个动作向量（包含关节角度、末端执行器位置等信息）。

为了应用Flow Matching，SmolVLA将真实动作序列进行噪声化：

$Atτ=τϵ+(1−τ)AtA_t^\tau = \tau \epsilon + (1-\tau) A_t$

让我们详细解释这个公式：

$A_t$ ：真实的动作序列（我们的目标）
$ϵ∼N(0,I)\epsilon \sim \mathcal{N}(0, I)$ ：从标准正态分布采样的噪声序列
$τ∈[0,1]\tau \in [0,1]$ ：噪声水平参数，控制噪声和真实动作的混合比例
$AtτA_t^\tau$ ：噪声化后的动作序列

这个公式的含义：

当 $τ=0\tau = 0$ 时： $At0=0⋅ϵ+1⋅At=AtA_t^0 = 0 \cdot \epsilon + 1 \cdot A_t = A_t$ （纯真实动作）
当 $τ=1\tau = 1$ 时： $At1=1⋅ϵ+0⋅At=ϵA_t^1 = 1 \cdot \epsilon + 0 \cdot A_t = \epsilon$ （纯噪声）
当 $τ=0.3\tau = 0.3$ 时： $At0.3=0.3ϵ+0.7AtA_t^{0.3} = 0.3 \epsilon + 0.7 A_t$ （30%噪声 + 70%真实动作）

对应的条件向量场为：

$u(Atτ∣At)=ϵ−Atu(A_t^\tau | A_t) = \epsilon - A_t$

这个向量场的含义是：无论当前处于哪个噪声水平，移动方向都是从噪声指向真实动作。这是一个非常直观的结果：我们希望从随机的动作逐渐变成有意义的动作。

SmolVLA的损失函数

根据论文，SmolVLA的训练损失函数为：

$Lτ(θ)=Ep(At∣ot),q(Atτ∣At)[∥vθ(Atτ,ot)−u(Atτ∣At)∥2]\mathcal{L}^\tau(\theta) = \mathbb{E}_{p(A_t|o_t), q(A_t^\tau|A_t)}[\|v_\theta(A_t^\tau, o_t) - u(A_t^\tau | A_t)\|^2]$

让我们逐项解释这个公式：

$Lτ(θ)\mathcal{L}^\tau(\theta)$ ：损失函数，上标 $τ\tau$ 表示这是针对特定噪声水平的损失
$p(A_t|o_t)$ ：给定观察 $o_t$ （包括图像、语言、状态）时真实动作序列 $A_t$ 的分布
$q(Atτ∣At)q(A_t^\tau|A_t)$ ：给定真实动作 $A_t$ 时噪声动作 $AtτA_t^\tau$ 的分布
$vθ(Atτ,ot)v_\theta(A_t^\tau, o_t)$ ：Action Expert网络预测的向量场
$u(Atτ∣At)=ϵ−Atu(A_t^\tau | A_t) = \epsilon - A_t$ ：真实的条件向量场

这个损失函数的训练过程：

从数据集中采样一个观察-动作对 $o_t, A_t)$
采样噪声 $ϵ\epsilon$ 和噪声水平 $τ\tau$
计算噪声动作 $Atτ=τϵ+(1−τ)AtA_t^\tau = \tau \epsilon + (1-\tau) A_t$
计算真实向量场 $u(Atτ∣At)=ϵ−Atu(A_t^\tau | A_t) = \epsilon - A_t$
让网络预测向量场 $vθ(Atτ,ot)v_\theta(A_t^\tau, o_t)$
计算预测值和真实值之间的MSE损失

现在我们已经理解了SmolVLA如何将Flow Matching理论应用到动作生成中，接下来让我们深入分析Action Expert的具体实现，看看它是如何将理论转化为可执行的代码的。

四、Action Expert的核心架构设计

交替注意力机制：创新的网络结构

SmolVLA的Action Expert最重要的创新之一是采用了交替注意力机制。为了理解这个设计的巧妙之处，让我们先回顾一下传统Transformer的结构。

在标准的Transformer解码器中，每一层都包含两种注意力机制：

自注意力(Self-Attention)：让序列中的每个元素关注序列中的其他元素
交叉注意力(Cross-Attention)：让序列中的元素关注来自编码器的信息

而SmolVLA采用了一种更轻量级的设计：每一层只包含一种注意力机制，交替使用自注意力和交叉注意力层。

这种设计的具体含义是：

交叉注意力层：动作token与VLM特征进行交互，从图像和语言中获取任务相关信息
自注意力层：动作token之间相互交互，确保生成的动作序列在时间上的连贯性

让我们用一个具体的例子来理解这个过程。假设机器人需要执行"拿起红色杯子"这个任务：

交叉注意力层：动作token会关注VLM提取的特征，比如"红色"、“杯子”、"拿起"等语义信息，以及图像中杯子的位置信息
自注意力层：动作序列中的每个动作会关注之前的动作，确保整个动作序列是合理的（比如先伸手，再抓取，最后提起）

因果掩码：保证时序逻辑

在自注意力层中，SmolVLA使用了因果掩码(Causal Mask)，这是一个非常重要的设计细节。让我们通过代码来理解：

# 创建因果掩码，确保时序一致性
att_masks += [1] * self.config.chunk_size  # 动作token使用因果掩码

因果掩码的作用是确保每个动作token只能关注到它之前的动作token，而不能"看到"未来的动作。这个设计有三个重要目的：

防止信息泄露：在训练时，如果允许当前动作关注未来动作，模型可能会"作弊"，利用未来信息来预测当前动作，这在实际应用中是不可能的。

保持时序一致性：机器人的动作必须按照时间顺序执行，当前动作只能基于过去和现在的信息来决定。

支持自回归生成：在推理时，模型可以逐步生成动作序列，每次生成一个动作，然后基于已生成的动作来预测下一个动作。

多模态特征的投影和融合

SmolVLA需要处理来自不同模态的信息：图像、语言、机器人状态和动作。为了让这些不同类型的信息能够有效交互，系统使用了多个线性投影层。

状态投影：连接机器人与VLM

self.state_proj = nn.Linear(
    self.config.max_state_dim, 
    self.vlm_with_expert.config.text_config.hidden_size
)

机器人的状态信息（如关节角度、末端执行器位置、夹爪开合状态等）通常是数值型的，而VLM处理的是高维的语义特征。状态投影层的作用是将这些数值信息转换到VLM的特征空间，使其能够与视觉和语言特征进行交互。

举个例子：如果机器人当前的夹爪是张开的，这个状态信息经过投影后，可以与语言指令中的"抓取"动作产生关联，帮助模型理解当前状态与目标动作的关系。

动作投影：连接动作空间与特征空间

# 动作输入投影：将动作投影到Expert特征空间
self.action_in_proj = nn.Linear(
    self.config.max_action_dim, 
    self.vlm_with_expert.expert_hidden_size
)

# 动作输出投影：将Expert输出投影回动作空间
self.action_out_proj = nn.Linear(
    self.vlm_with_expert.expert_hidden_size, 
    self.config.max_action_dim
)

动作投影包括两个方向：

输入投影：将原始动作空间（比如7维的机械臂关节角度）投影到Action Expert的高维特征空间，使其能够与其他特征进行复杂的交互。

输出投影：将Action Expert处理后的高维特征投影回原始动作空间，得到具体的动作指令。

这种设计的好处是，在高维特征空间中，模型可以进行更复杂的推理和计算，而最终输出仍然是机器人可以直接执行的动作指令。

时间嵌入：编码连续时间信息

时间信息在Flow Matching中至关重要，SmolVLA使用正弦余弦位置编码来表示时间：

def create_sinusoidal_pos_embedding(timestep, hidden_size):
    half_dim = hidden_size // 2
    emb = math.log(10000) / (half_dim - 1)
    emb = torch.exp(torch.arange(half_dim) * -emb)
    emb = timestep[:, None] * emb[None, :]
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
    return emb

让我们详细解释这个函数：

half_dim = hidden_size // 2：将嵌入维度分成两半，一半用于正弦函数，一半用于余弦函数。

emb = math.log(10000) / (half_dim - 1)：计算频率的基础值，10000是一个经验值，用于控制不同维度的频率范围。

emb = torch.exp(torch.arange(half_dim) * -emb)：为每个维度计算不同的频率，低维度对应低频率（捕捉长期模式），高维度对应高频率（捕捉短期模式）。

emb = timestep[:, None] * emb[None, :]：将时间值与频率相乘，得到每个维度的相位。

torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)：使用正弦和余弦函数生成最终的嵌入，这样可以唯一地表示任何时间值。

这种编码方式的优势是：

能够表示连续的时间值（不像离散的位置编码）
对不同的时间尺度都有良好的表示能力
具有良好的插值性质，相近的时间值会有相似的编码

五、训练过程：理论与实践的结合

现在让我们看看SmolVLA是如何在训练过程中实现Flow Matching算法的。

输入处理：前缀与后缀的分离

SmolVLA采用了一种巧妙的输入处理策略，将所有输入分为两部分：

前缀(Prefix)：包含任务相关的上下文信息

图像：当前环境的视觉信息
语言：用户的指令文本
状态：机器人当前的状态

后缀(Suffix)：包含Flow Matching相关的信息

噪声动作：经过噪声化的动作序列
时间：Flow Matching中的时间参数

这种分离的好处是可以针对不同类型的信息采用不同的处理策略。

前缀嵌入：构建任务上下文

def embed_prefix(self, images, img_masks, lang_tokens, lang_masks, state):
    """嵌入前缀：图像、语言、状态信息"""
    embs = []
    pad_masks = []
    att_masks = []
    
    # 处理图像嵌入
    if images is not None:
        img_emb = self.vlm_with_expert.embed_image(images)
        embs.append(img_emb)
        pad_masks.append(img_masks)
        att_masks += [0] * img_emb.shape[1]  # 图像token可以被后续token关注
    
    # 处理语言嵌入
    lang_emb = self.vlm_with_expert.embed_language_tokens(lang_tokens)
    lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])  # 缩放嵌入
    embs.append(lang_emb)
    pad_masks.append(lang_masks)
    att_masks += [0] * lang_emb.shape[1]
    
    # 处理状态嵌入
    state_emb = self.state_proj(state)
    if state_emb.ndim == 2:
        state_emb = state_emb[:, None, :]  # 添加序列维度
    embs.append(state_emb)
    
    # 状态mask设置
    bsize, states_seq_len = state_emb.shape[:2]
    state_mask = torch.ones(bsize, states_seq_len, dtype=torch.bool, device=state_emb.device)
    pad_masks.append(state_mask)
    att_masks += [1] * states_seq_len  # 状态token不能被动作token关注
    
    # 拼接所有嵌入
    embs = torch.cat(embs, dim=1)
    pad_masks = torch.cat(pad_masks, dim=1)
    att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
    
    return embs, pad_masks, att_masks

这个函数中有几个重要的设计细节：

嵌入缩放：lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
这是Transformer中的标准做法，用于平衡不同输入的贡献。缩放因子 $d\sqrt{d}$ 可以防止嵌入值过大，影响训练稳定性。

注意力掩码的精心设计：

图像和语言token设置为att_masks = 0，表示可以被动作token关注
状态token设置为att_masks = 1，表示不能被动作token关注

这种设计的深层原理是：

图像和语言提供任务的上下文信息，应该被动作生成过程参考
状态信息是当前时刻的即时信息，如果允许未来的动作关注当前状态，可能会导致信息泄露

后缀嵌入：融合动作与时间

def embed_suffix(self, noisy_actions, timestep):
    """嵌入后缀：噪声动作和时间信息"""
    # 动作嵌入
    action_emb = self.action_in_proj(noisy_actions)
    
    # 时间嵌入：使用正弦余弦位置编码
    time_emb = create_sinusoidal_pos_embedding(
        timestep, self.vlm_with_expert.expert_hidden_size, 
        self.config.min_period, self.config.max_period, device=device
    )
    
    # 扩展时间嵌入到与动作嵌入相同的形状
    time_emb = time_emb[:, None, :].expand_as(action_emb)
    
    # 拼接动作和时间嵌入
    action_time_emb = torch.cat([action_emb, time_emb], dim=2)
    
    # 通过MLP融合
    action_time_emb = self.action_time_mlp_in(action_time_emb)
    action_time_emb = F.silu(action_time_emb)  # SiLU激活函数
    action_time_emb = self.action_time_mlp_out(action_time_emb)
    
    return action_time_emb, pad_masks, att_masks

时间嵌入的扩展：时间是一个标量值，但动作是一个序列。通过expand_as操作，时间信息被复制到动作序列的每个位置，确保每个动作都能获得时间信息。

特征融合策略：使用拼接(concatenation)而不是相加(addition)来融合动作和时间特征，这样可以保留更多的信息。然后通过MLP进一步处理，让模型学习如何最好地结合这两种信息。

训练的核心循环

训练时的前向传播完整实现了Flow Matching算法：

def forward(self, images, img_masks, lang_tokens, lang_masks, state, actions, noise=None, time=None):
    """训练时的前向传播"""
    # 步骤1：采样噪声和时间
    if noise is None:
        noise = self.sample_noise(actions.shape, actions.device)
    if time is None:
        time = self.sample_time(actions.shape[0], actions.device)
    
    # 步骤2：计算噪声动作
    time_expanded = time[:, None, None]
    x_t = time_expanded * noise + (1 - time_expanded) * actions
    
    # 步骤3：计算目标向量场
    u_t = noise - actions
    
    # 步骤4：嵌入处理
    prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
        images, img_masks, lang_tokens, lang_masks, state=state
    )
    suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, time)
    
    # 步骤5：拼接和网络处理
    pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
    att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
    
    att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
    position_ids = torch.cumsum(pad_masks, dim=1) - 1
    
    (_, suffix_out), _ = self.vlm_with_expert.forward(
        attention_mask=att_2d_masks,
        position_ids=position_ids,
        past_key_values=None,
        inputs_embeds=[prefix_embs, suffix_embs],
        use_cache=False,
        fill_kv_cache=False,
    )
    
    # 步骤6：预测向量场
    suffix_out = suffix_out[:, -self.config.chunk_size:]
    v_t = self.action_out_proj(suffix_out.to(dtype=torch.float32))
    
    # 步骤7：计算损失
    losses = F.mse_loss(u_t, v_t, reduction="none")
    return losses

让我们逐步分析这个训练过程：

步骤1-2：噪声化过程
这里实现了公式 $Atτ=τϵ+(1−τ)AtA_t^\tau = \tau \epsilon + (1-\tau) A_t$ 。time_expanded的维度调整确保了正确的广播操作。

步骤3：目标计算
计算真实的条件向量场 $u(Atτ∣At)=ϵ−Atu(A_t^\tau | A_t) = \epsilon - A_t$ ，这是模型需要学习预测的目标。

步骤4-5：特征处理
将多模态输入转换为网络可以处理的统一表示，并设置正确的注意力掩码。

步骤6：网络预测
通过Action Expert网络预测向量场 $vθ(Atτ,ot)v_\theta(A_t^\tau, o_t)$ 。

步骤7：损失计算
使用MSE损失衡量预测向量场与真实向量场的差距。

六、推理过程：从噪声到动作的生成

推理阶段是Flow Matching理论的真正体现，模型需要从纯噪声开始，逐步生成有意义的动作序列。

ODE求解的数值实现

def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state, noise=None):
    """推理时的动作采样"""
    # 初始化：从纯噪声开始
    if noise is None:
        actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
        noise = self.sample_noise(actions_shape, device)
    
    # 优化：预计算前缀的KV缓存
    prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
        images, img_masks, lang_tokens, lang_masks, state=state
    )
    _, past_key_values = self.vlm_with_expert.forward(
        # ... 预计算KV缓存
    )
    
    # ODE求解：从τ=1（纯噪声）到τ=0（真实动作）
    dt = -1.0 / self.config.num_steps
    x_t = noise
    time = torch.tensor(1.0, dtype=torch.float32, device=device)
    
    # Euler方法逐步求解
    while time >= -dt / 2:
        expanded_time = time.expand(bsize)
        v_t = self.denoise_step(prefix_pad_masks, past_key_values, x_t, expanded_time)
        
        # Euler步骤：x_{t+dt} = x_t + dt * v_t
        x_t += dt * v_t
        time += dt
    
    return x_t

这个推理过程有几个关键点：

从噪声到数据的路径：我们从 $τ=1\tau=1$ （纯噪声）开始，逐步移动到 $τ=0\tau=0$ （真实数据）。负的步长dt = -1.0 / self.config.num_steps确保了这个方向。

Euler方法：这是求解ODE的最简单方法。虽然存在更精确的方法（如Runge-Kutta），但Euler方法在实践中已经足够，且计算效率更高。

KV缓存优化：由于前缀部分在整个采样过程中不变，预先计算并缓存其键值对可以显著提高推理速度。

单步去噪的实现

每个时间步的核心是预测当前状态下的向量场：

def denoise_step(self, prefix_pad_masks, past_key_values, x_t, timestep):
    """单步去噪过程"""
    # 嵌入当前状态和时间
    suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, timestep)
    
    # 构造完整的输入
    pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
    att_masks = torch.cat([torch.zeros_like(prefix_pad_masks), suffix_att_masks], dim=1)
    
    # 通过网络预测向量场
    (_, suffix_out), _ = self.vlm_with_expert.forward(
        attention_mask=make_att_2d_masks(pad_masks, att_masks),
        position_ids=torch.cumsum(pad_masks, dim=1) - 1,
        past_key_values=past_key_values,
        inputs_embeds=[None, suffix_embs],
        use_cache=self.config.use_cache,
        fill_kv_cache=False,
    )
    
    # 提取并投影输出
    suffix_out = suffix_out[:, -self.config.chunk_size:]
    v_t = self.action_out_proj(suffix_out.to(dtype=torch.float32))
    
    return v_t

这个函数的核心思想是：给定当前的噪声动作和时间，预测应该朝哪个方向"移动"来逐渐接近真实的动作序列。通过多次迭代，最终从纯噪声生成有意义的动作序列。

整个推理过程可以理解为一个"去噪"过程：我们从一个完全随机的动作序列开始，然后在每个时间步，网络告诉我们如何调整这个序列，使其更接近真实的、有意义的动作序列。经过多次迭代后，我们得到了一个可以让机器人执行的动作序列。

总结

SmolVLA中的Flow Matching Action Expert代表了机器人动作生成领域的一个重要进展。通过巧妙地将Flow Matching理论应用到动作生成任务中，它实现了理论创新与工程实践的完美结合。Flow Matching的无仿真训练方法大大提高了训练效率，同时保持了生成质量；交替使用自注意力和交叉注意力的架构设计既保证了计算效率，又维持了模型的表达能力；模块化的设计使得系统可以灵活地适应不同的机器人平台和任务需求。从数学理论到代码实现，每一个环节都体现了深思熟虑的设计选择，这些选择共同构成了一个高效、稳定、实用的机器人动作生成系统，为机器人学习领域提供了新的研究方向和实践范式。

参考文献

[1] Shukor, M., Aubakirova, D., Capuano, F., et al. (2025). SmolVLA: A vision-language-action model for affordable and efficient robotics. arXiv preprint arXiv:2506.01844.

[2] Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2023). Flow Matching for Generative Modeling. arXiv preprint arXiv:2210.02747. https://arxiv.org/abs/2210.02747

[3] Fjelde, T., Mathieu, E., & Dutordoir, V. (2024). An introduction to Flow Matching. Cambridge MLG Blog. https://mlg.eng.cam.ac.uk/blog/2024/01/20/flow-matching.html

[4] Chen, R. T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. (2018). Neural Ordinary Differential Equations. Advances in Neural Information Processing Systems, 31.

[5] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. Advances in Neural Information Processing Systems, 33, 6840-6851.

深入理解SmolVLA中的Flow Matching Action Expert：从理论到实现

引言