Flow Matching 入门详解：从扩散模型到连续动作生成（含代码讲解）

class BagelConfig(PretrainedConfig):
    def __init__(
        self,
        visual_gen=True,        # 开启 Flow Matching 生成
        visual_und=True,        # 开启视觉理解
        llm_config=None,        # Qwen2 config
        vit_config=None,        # SIGLIP config
        vae_config=None,        # VAE config（输出 latent）
        latent_patch_size=2,    # 每个 latent patch 的尺寸
        max_latent_size=32,     # 最大 latent 网格
        timestep_shift=1.0,     # timestep 偏移（推理时常用 3.0）
        ...
    )

# -----------------------------------------------------------
# Step 1: 从 VAE latent 构造干净的 packed_latent
# -----------------------------------------------------------
p = self.latent_patch_size          # 通常为 2
packed_latent = []
for latent, (h, w) in zip(padded_latent, patchified_vae_latent_shapes):
    # VAE latent shape: [C, H, W] → reshape to patches
    # 把 latent 切成 p×p 的 patch，每个 patch 展平成向量
    latent = latent[:, :h * p, :w * p].reshape(self.latent_channel, h, p, w, p)
    latent = torch.einsum("chpwq->hwpqc", latent).reshape(-1, p * p * self.latent_channel)
    packed_latent.append(latent)
packed_latent_clean = torch.cat(packed_latent, dim=0)

# -----------------------------------------------------------
# Step 2: 采样噪声 + 线性插值（CFM 核心）
# -----------------------------------------------------------
noise = torch.randn_like(packed_latent_clean)          # x₀: 高斯噪声
packed_timesteps = torch.sigmoid(packed_timesteps)     # t: 0→1 的 timestep
# timestep shift: 让模型在 t 接近 0 时"看得更细"
packed_timesteps = self.timestep_shift * packed_timesteps / \
                   (1 + (self.timestep_shift - 1) * packed_timesteps)
# x_t = (1-t) * x₀ + t * x₁  直线插值
packed_latent = (1 - packed_timesteps[:, None]) * packed_latent_clean \
                + packed_timesteps[:, None] * noise

# -----------------------------------------------------------
# Step 3: 把 latent patch 映射到 LLM 隐空间 + 加 timestep/位置编码
# -----------------------------------------------------------
packed_timestep_embeds = self.time_embedder(packed_timesteps)
latent_token_pos_emb = self.latent_pos_embed(packed_latent_position_ids)
# vae2llm: Linear 把 patch 投影到 hidden_size
packed_latent = self.vae2llm(packed_latent) \
                + packed_timestep_embeds + latent_token_pos_emb
# 插入到 packed_sequence 的对应位置（与文本 token 拼接）
packed_sequence[packed_vae_token_indexes] = packed_latent

# -----------------------------------------------------------
# Step 4: 过 LLM backbone，得到隐状态
# -----------------------------------------------------------
last_hidden_state = self.language_model(packed_sequence=packed_sequence, ...)

# -----------------------------------------------------------
# Step 5: 预测速度场 + MSE loss
# -----------------------------------------------------------
packed_mse_preds = self.llm2vae(last_hidden_state[mse_loss_indexes])
# 注意这里的 target 方向：x₁ - x₀  即"从数据指向噪声"
# （与论文中 u_t = x₁ - x₀ 一致，符号取决于路径定义方向）
target = noise - packed_latent_clean
has_mse = packed_timesteps > 0
mse = (packed_mse_preds - target[has_mse]) ** 2

class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) *
            torch.arange(start=0, end=half, dtype=torch.float32) / half
        ).to(device=t.device)
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        t_emb = self.mlp(t_freq)          # 256 → hidden_size
        return t_emb

def _sde_step_with_logprob(self, model_output, timestep, prev_timestep,
                           d_timestep, sample, prev_sample=None,
                           noise_level=0.8):
    model_output = model_output.float()
    sample = sample.float()

    # 计算当前步的噪声标准差
    # std_dev_t = sqrt(t / (1 - t')) * noise_level
    # 其中 t' = max(sigma_max, t) 防止除零
    std_dev_t = torch.sqrt(
        timestep / (1 - torch.where(timestep == 1, sigma_max, timestep))
    ) * noise_level

    # 预测下一步均值（Euler 步 + 噪声修正项）
    #   x_{t-1} = x_t*(1 + σ²/(2t)*dt) + v_t*(1 + σ²*(1-t)/(2t))*dt
    # 当 noise_level=0 时退化为纯 Euler: x_{t-1} = x_t + v_t * dt
    prev_sample_mean = (
        sample * (1 + std_dev_t**2 / (2 * timestep) * d_timestep)
        + model_output * (1 + std_dev_t**2 * (1 - timestep) / (2 * timestep)) * d_timestep
    )

    if prev_sample is None:
        variance_noise = randn_tensor(model_output.shape, ...)
        prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-d_timestep) * variance_noise

    # log_prob 用于 RL 策略梯度（Flow-GRPO 需要）
    log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) \
               / (2 * (std_dev_t * torch.sqrt(-d_timestep))**2)
    log_prob = log_prob.mean()

    return prev_sample, log_prob, prev_sample_mean, std_dev_t

def generate_image(self, ...):
    x_t = packed_init_noises                      # 起点：纯噪声

    # 准备离散 timestep 序列（从 1 → 0）
    timesteps = torch.linspace(1, 0, num_timesteps)  # 例如 50 步
    timesteps = timestep_shift * timesteps / \
                (1 + (timestep_shift - 1) * timesteps)
    dts = timesteps[1:] - timesteps[:-1]

    for i, t in enumerate(timesteps[:-1]):
        # -----------------------------------------------------------
        # 预测速度场（含可选的 CFG）
        # -----------------------------------------------------------
        v_t = self._forward_flow(
            x_t=x_t, timestep=torch.tensor([t]),
            cfg_text_scale=cfg_text_scale,   # 文本 CFG 强度
            cfg_img_scale=cfg_img_scale,     # 图像 CFG 强度
            ...
        )
        # -----------------------------------------------------------
        # SDE step（含可选的噪声注入）
        # -----------------------------------------------------------
        x_t, log_prob, _, _ = self._sde_step_with_logprob(
            v_t, timesteps[i], timesteps[i+1], dts[i], x_t,
            noise_level=cur_noise_level
        )

    return x_t  # 生成结果

def _forward_flow(self, x_t, timestep, ..., cfg_text_scale, cfg_img_scale):
    # 1) 条件前向（有 text/image 条件）
    v_t = self.llm2vae(llm_output)

    # 2) 无条件前向（text CFG：用空文本做条件）
    if cfg_text_scale > 1.0:
        cfg_text_v_t = self.llm2vae(cfg_text_llm_output)

    # 3) 图像无条件前向（image CFG）
    if cfg_img_scale > 1.0:
        cfg_img_v_t = self.llm2vae(cfg_img_llm_output)

    # 4) CFG 插值 + renormalization（防止 CFG 导致向量范数爆炸）
    #    v = v_uncond + scale * (v_cond - v_uncond)
    v_t = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t)
    v_t = v_t * (norm(v_t) / norm(v_t_))  # CFG-Renorm 保持速度尺度

    return v_t

# transition: x_{t+1} ~ N(mean_t, σ_t²)
# log p(x_{t+1} | x_t, θ) = -(x_{t+1} - mean_t)² / (2σ_t²)
log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2) \
           / (2 * (std_dev_t * torch.sqrt(-d_timestep))**2)
log_prob = log_prob.mean()

# generate_image() 中保存每一步的 latent 和 log_prob
all_latents = []     # 每一步的 latent
all_log_probs = []   # 每一步的对数概率
all_timesteps = []   # 对应的时间步

for i, t in enumerate(timesteps):
    v_t = self._forward_flow(x_t, ...)
    x_t, log_prob, _, _ = self._sde_step_with_logprob(v_t, ...)
    # 在指定窗口内保存中间结果用于 RL 训练
    if i >= sde_timestep_begin and i < sde_timestep_begin + window_size:
        all_latents.append(x_t)
        all_log_probs.append(log_prob)
        all_timesteps.append(t)

def generate_image_learn(self, sample, grpo_config, accelerator, optimizer, ...):
    latents = sample["latents"]           # 采样时保存的 latent 序列
    prev_latents = sample["prev_latents"] # 上一步的 latent
    timesteps = sample["timesteps"]
    advantages = torch.clamp(
        sample["advantages"],
        -grpo_config.train.adv_clip_max,
        grpo_config.train.adv_clip_max,
    )

    for i, t in enumerate(timesteps):
        with accelerator.accumulate(transformer):
            # 用当前策略预测速度场（online）
            v_t = self._forward_flow(x_t=latents[i], timestep=t, ...)

            # 计算 log_prob（含重参数化）
            _, log_prob, prev_sample_mean, std_dev_t = self._sde_step_with_logprob(
                v_t, timesteps[i], timesteps[i+1], dts[i],
                latents[i], prev_sample=prev_latents[i], ...
            )

            # KL 正则：对 reference model 也做一步推理
            if grpo_config.train.beta > 0:
                v_t_ref = self._forward_flow(..., ref_model=True)
                _, _, prev_sample_mean_ref, _ = self._sde_step_with_logprob(
                    v_t_ref, ...
                )

            # -----------------------------------------------------------
            # GRPO 策略梯度（核心 4 行）
            # -----------------------------------------------------------
            ratio = torch.exp(log_prob - sample["log_probs"][i])
            unclipped_loss = -advantages * ratio
            clipped_loss = -advantages * torch.clamp(
                ratio,
                1.0 - grpo_config.train.clip_range_lt,
                1.0 + grpo_config.train.clip_range_gt,
            )
            policy_loss = torch.mean(torch.maximum(unclipped_loss, clipped_loss))

            # KL 散度（高斯分布的 KL 有闭式解）
            if grpo_config.train.beta > 0:
                kl_loss = ((prev_sample_mean - prev_sample_mean_ref) ** 2).mean() \
                          / (2 * std_dev_t ** 2)
                loss = policy_loss + grpo_config.train.beta * kl_loss
            else:
                loss = policy_loss

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