【具身笔记】【原理探究】二、Flow Matching(流匹配)代码实践

m0_75114363

37人浏览 · 2026-10-27 00:48:14

m0_75114363 · 2026-10-27 00:48:14 发布

概述

本系列为我在vla 上的同步实践笔记，之前的实践中我大多都是在调用指令、修改参数，没有去深入了解相关的底层原理，因此在本系列中我将尽可能地探究底层原理。本文将展示我通过代码小实验对flow matching的理解。

1. 实验设计

本次实验的输入为带有噪声信号的二维轨迹，仿照pi0.5的流程，通过流匹配求出每一轮的轨迹变化速度，10次积分后得到最终轨迹。

2. 代码

import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import torch
from torch import nn

N_POINTS = 50
TRAIN_STEPS = 20000
BATCH_SIZE = 128
GENERATE_STEPS = 10

class FlowMLP(nn.Module):
    def __init__(self):
        super().__init__()
        # 输入：50×2个动作数值 + 时间t + 曲线幅度条件。即初始是50个2维的点坐标
        # 输出：50×2个速度数值。
        self.net = nn.Sequential(
            # 加的2为：时间t一维标量 + 曲线幅度1维标量
            nn.Linear(N_POINTS * 2 + 2, 256),
            nn.SiLU(),
            nn.Linear(256, 256),
            nn.SiLU(),
            nn.Linear(256, N_POINTS * 2)
        )

    def forward(self, x_t, t, condition):
        # xt的第0维(batch维)不变，后面的维度拉平，变成[x1, y1, x2, y2, ..., x50, y50]
        # 拼接时，也是batch维不变，后面的拼好，变成[x1, y1, x2, y2, x3, y3, ..., x50, y50, t, condition]
        inputs = torch.cat(
            [x_t.flatten(start_dim=1), t, condition], dim=1
        )
        return self.net(inputs).reshape_as(x_t)

# 下面这个函数用来生成标准答案曲线
def make_trajectory(condition):
    # 每个幅度条件对应一条完整曲线：(batch, 50, 2)。
    # 第一步，生成x坐标
    # linspace从 -1 到 1，均匀取 50 个点
    # expand是为了增加batch维，变成(B, 50)
    x = torch.linspace(
        -1, 1, N_POINTS, device=condition.device
    ).expand(condition.shape[0], -1)
    # 第二步，生成y坐标，即正弦函数，高度由condition决定
    y = condition * torch.sin(torch.pi * x)
    # 把生成的x和y一一配对，变成(B, 50, 2)，即B批，50对，每对都是(x,y)
    return torch.stack([x, y], dim=-1)

def main():
    torch.manual_seed(1000)
    device = torch.device(
        "mps" if torch.backends.mps.is_available() else "cpu"
    )

    model = FlowMLP().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    print("设备：", device)
    print("参数量：", sum(p.numel() for p in model.parameters()))

    # 1.训练
    model.train()
    losses = []

    for step in range(1, TRAIN_STEPS + 1):
        # 幅度在0.6～1.4之间变化，作为网络的条件输入。
        condition = 0.6 + 0.8 * torch.rand(
            BATCH_SIZE, 1, device=device
        )
        # 生成标准答案。由于每一轮训练时的condition随机，故标准答案也随机，由此实现泛化性
        actions = make_trajectory(condition)
        # 噪声
        noise = torch.rand_like(actions)
        # 随机时间步
        t = 0.001 + 0.998 * torch.rand(
            BATCH_SIZE, 1, device=device
        )

        # 与π0.5采用相同的线性加噪和速度监督公式。
        # 每一step的训练中，time都不一样，因此该step输入的带噪时间步也都不一样，模型正是在多种带噪时间步的输入下一点一点学会的
        time = t[:, :, None]
        # 初始的带噪时间步
        x_t = (1 - time) * actions + time * noise
        # 参考答案
        target_velocity = noise - actions

        # 输入为这一轮训练的带噪时间步 + 这一轮对应的t参数 + 这一轮对应的condition参数
        predicted_velocity = model(x_t, t, condition)
        loss = nn.functional.mse_loss(
            predicted_velocity, target_velocity
        )

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        value = loss.item()
        losses.append(value)

        if step == 1:
            print("真实轨迹形状：", tuple(actions.shape))
            print("带噪轨迹形状：", tuple(x_t.shape))
            print("预测速度形状：", tuple(predicted_velocity.shape))

        if step == 1 or step % 500 == 0:
            print(f"step={step}, loss={value}")

    # 2. 推理
    model.eval()
    snapshots = {}

    with torch.inference_mode():
        # 生成随机的初始轨迹
        x_t = torch.randn(1, N_POINTS, 2, device=device)
        snapshots[0] = x_t[0].cpu().numpy().copy()

        # 标准参考答案
        condition = torch.tensor([[1.0]], device=device)
        reference = make_trajectory(condition)
        initial_mse = nn.functional.mse_loss(
            x_t, reference
        ).item()

        # 用于速度积分
        dt = -1.0 / GENERATE_STEPS

        for step in range(GENERATE_STEPS):
            t = torch.full(
                (1, 1), 1 + step * dt, device=device
            )
            velocity = model(x_t, t, condition)
            # 积分：当前带噪动作 + 当前这一步的速度 * 时间
            x_t = x_t + dt * velocity

            if step + 1 in (2, 5, 8, 10):
                snapshots[step + 1] = (
                    x_t[0].cpu().numpy().copy()
                )

        # 最后生成的loss
        generated_mse = nn.functional.mse_loss(
            x_t, reference
        ).item()
        # 参考轨迹
        reference_points = reference[0].cpu().numpy()

        assert torch.isfinite(x_t).all().item()

    # 3. 画图
    # 3. 一张图：训练loss + 噪声到曲线的生成过程。
    fig, axes = plt.subplots(2, 3, figsize=(15, 9))
    axes = axes.flatten()

    axes[0].plot(losses, linewidth=1)
    axes[0].set_yscale("log")
    axes[0].set_title("Training velocity MSE")
    axes[0].set_xlabel("Training step")
    axes[0].set_ylabel("Loss")
    axes[0].grid(alpha=0.25)

    limit = max(
        1.5, max(abs(points).max() for points in snapshots.values())
    ) * 1.05

    for ax, (step, points) in zip(axes[1:], snapshots.items()):
        ax.plot(
            reference_points[:, 0],
            reference_points[:, 1],
            "--",
            color="gray",
            label="Target",
        )
        ax.plot(
            points[:, 0], points[:, 1],
            "o-", markersize=3, linewidth=1, label="Generated",
        )
        ax.set_title(
            f"After {step} updates, t={1-step/GENERATE_STEPS:.1f}"
        )
        ax.set_xlim(-limit, limit)
        ax.set_ylim(-limit, limit)
        ax.set_aspect("equal")
        ax.grid(alpha=0.25)
        ax.legend(fontsize=8)

    fig.suptitle(
        "Flow Matching: 50 two-dimensional points form one trajectory"
    )
    fig.tight_layout()

    output_dir = (
        Path(__file__).resolve().parents[1]
        / "outputs/week6/c1_flow_matching"
    )
    output_dir.mkdir(parents=True, exist_ok=True)

    figure_path = output_dir / "c1_flow_matching.png"
    fig.savefig(figure_path, dpi=180)
    plt.close(fig)

    metrics = {
        "device": str(device),
        "train_steps": TRAIN_STEPS,
        "generation_steps": GENERATE_STEPS,
        "training_losses": losses,
        "initial_noise_mse": initial_mse,
        "generated_trajectory_mse": generated_mse,
    }
    (output_dir / "metrics.json").write_text(
        json.dumps(metrics, indent=2)
    )

    print("初始噪声与目标轨迹的MSE：", initial_mse)
    print("生成轨迹与目标轨迹的MSE：", generated_mse)
    print("生成过程图：", figure_path)

    assert generated_mse < initial_mse, "生成未改善，请保留输出排查"
    print("FLOW_MATCHING_DEMO_CHECK_PASS")

if __name__ == "__main__":
    main()

3. 结果

参数量： 117860
真实轨迹形状： (128, 50, 2)
带噪轨迹形状： (128, 50, 2)
预测速度形状： (128, 50, 2)
step=1, loss=0.7628627419471741
step=500, loss=0.056497376412153244
step=1000, loss=0.03080480545759201
step=1500, loss=0.017268814146518707
step=2000, loss=0.017937956377863884
step=2500, loss=0.013999688439071178
step=3000, loss=0.014280745759606361
step=3500, loss=0.0105742784217
step=4000, loss=0.01215374656021595
初始噪声与目标轨迹的MSE： 1.0767486095428467
生成轨迹与目标轨迹的MSE： 0.10508836805820465
生成过程图： /Users/dukiya/Desktop/VLA-learning/pi0.5/outputs/week6/c1_flow_matching/c1_flow_matching.png
FLOW_MATCHING_DEMO_CHECK_PASS
(vla-dl) dukiya@MacBookPro pi0.5 % clear
(vla-dl) dukiya@MacBookPro pi0.5 % python -u "/Users/dukiya/Desktop/VLA-learning/pi0.5/scripts/C1_1_flow_matching.py"
设备： mps
参数量： 117860
真实轨迹形状： (128, 50, 2)
带噪轨迹形状： (128, 50, 2)
预测速度形状： (128, 50, 2)
step=1, loss=0.7628627419471741
step=500, loss=0.056497376412153244
step=1000, loss=0.03080480545759201
step=1500, loss=0.017268814146518707
step=2000, loss=0.017937956377863884
step=2500, loss=0.013999688439071178
step=3000, loss=0.014280745759606361
step=3500, loss=0.0105742784217
step=4000, loss=0.01215374656021595
step=4500, loss=0.010596787557005882
step=5000, loss=0.011056812480092049
step=5500, loss=0.011975659988820553
step=6000, loss=0.012458705343306065
step=6500, loss=0.00996104534715414
step=7000, loss=0.010937346145510674
step=7500, loss=0.007030962500721216
step=8000, loss=0.011274922639131546
step=8500, loss=0.009572314098477364
step=9000, loss=0.010419373400509357
step=9500, loss=0.00797924492508173
step=10000, loss=0.01026662066578865
step=10500, loss=0.009255047887563705
step=11000, loss=0.007038257550448179
step=11500, loss=0.010559530928730965
step=12000, loss=0.00735721318051219
step=12500, loss=0.00927362497895956
step=13000, loss=0.009580424055457115
step=13500, loss=0.0087526710703969
step=14000, loss=0.01084986887872219
step=14500, loss=0.01018508430570364
step=15000, loss=0.00851376261562109
step=15500, loss=0.009675636887550354
step=16000, loss=0.009681525640189648
step=16500, loss=0.00786056462675333
step=17000, loss=0.010169069282710552
step=17500, loss=0.009501800872385502
step=18000, loss=0.007643285673111677
step=18500, loss=0.00810040719807148
step=19000, loss=0.007944180630147457
step=19500, loss=0.007525602355599403
step=20000, loss=0.0067156539298594
初始噪声与目标轨迹的MSE： 1.3391687870025635
生成轨迹与目标轨迹的MSE： 0.0014069663593545556
生成过程图： /Users/dukiya/Desktop/VLA-learning/pi0.5/outputs/week6/c1_flow_matching/c1_flow_matching.png
FLOW_MATCHING_DEMO_CHECK_PASS
(vla-dl) dukiya@MacBookPro pi0.5 %