SimWAM 代码完整解读：从 Video DiT 到 Action DiT，从 SFT 到 FlowGRPO 的逐行拆解

# configs/model/simwam_navsim.yaml
# 这个文件定义了 SimWAM 的所有模型参数

# === Video Expert ===
model_id: "Wan-AI/Wan2.2-T2V-14B"     # 预训练视频模型
tokenizer_model_id: "Wan-AI/Wan2.2-T2V-14B"  # T5 文本编码器
tokenizer_max_len: 512                  # 文本最大长度
load_text_encoder: false                # 训练时不加载文本编码器（冻结）

# === Video DiT 配置 ===
video_dit_config:
  hidden_dim: 3072        # Video DiT 隐藏维度
  ffn_dim: 12288          # FFN 中间维度
  num_heads: 24           # 注意力头数
  attn_head_dim: 128      # 每头维度（3072/24=128）
  num_layers: 40          # Video DiT 层数
  freq_dim: 256           # 时间步 Fourier 编码维度

# === Action Expert 配置 ===
action_dit_config:
  action_dim: 3           # (x, y, heading) 三维度
  hidden_dim: 1024        # ActionDiT 隐藏维度
  ffn_dim: 4096           # FFN 中间维度
  num_heads: 24           # 注意力头数
  attn_head_dim: 128      # 每头维度
  num_layers: 30          # ActionDiT 层数（比 Video 少）
  text_dim: 4096          # 文本条件维度（与 T5 对齐）
  freq_dim: 256           # 时间步 Fourier 编码维度

# === 调度器配置 ===
video_scheduler:
  train_shift: 5.0        # 训练时的噪声偏移
  infer_shift: 5.0        # 推理时的噪声偏移
  num_train_timesteps: 1000  # 训练时间步数

action_scheduler:
  train_shift: 1.0        # ActionDiT 的训练噪声偏移
  infer_shift: 1.0        # ActionDiT 的推理噪声偏移
  num_train_timesteps: 10  # 10 步欧拉积分
  prediction_type: "velocity"  # 预测速度场
  sigma_clamp_min: 0.1    # SDE 噪声下限

# configs/task/navsim_uncond_front_384x672_1e-4.yaml
# SFT 训练的超参数

batch_size: 2              # 每 GPU 批大小
learning_rate: 1e-4        # 学习率
num_epochs: 100            # 训练轮数
lr_scheduler_type: "cosine"  # 余弦退火
weight_decay: 1e-2         # 权重衰减
gradient_accumulation_steps: 1  # 梯度累积

model:
  proprio_dim: 4           # 自车状态维度（速度、加速度等）
  loss:
    lambda_video: 1.0      # 视频损失权重
    lambda_action: 1.0     # 动作损失权重（等权重）

# configs/task/navsim_grpo_action_pdm_384x672_flowgrpo_lora.yaml
# FlowGRPO 训练的超参数

model:
  checkpoint_path: "path/to/sft/checkpoint.pt"  # SFT 权重

grpo:
  # === LoRA 配置 ===
  lora:
    enabled: true          # 只训练 LoRA 适配器
    r: 16                  # LoRA rank
    alpha: 32.0            # LoRA alpha（缩放系数 = alpha/r = 2）
    target_modules:        # 只在注意力层加 LoRA
      - q
      - k
      - v
      - o
  
  # === 采样配置 ===
  sample:
    sde_mode: rigorous     # 严格的 SDE 采样（score-corrected）
    noise_level: 0.1       # SDE 噪声级别
    num_samples: 8         # 每个场景采样 G=8 条轨迹
  
  # === 训练配置 ===
  train:
    ppo_clip_range: 0.02   # PPO clip 范围 ε
    adv_clip_max: 5.0      # advantage 截断到 ±5
    num_inner_epochs: 4    # 每批 rollout 复用 4 次
    learning_rate: 5e-6    # RL 学习率（比 SFT 小 20 倍）
  
  # === 奖励配置 ===
  reward:
    # NAVSIM PDM 奖励
    # PDMS = NC × DAC × (5·TTC + 5·EP + 2·C) / 12

# scripts/train_grpo.py（整个框架的入口）
import hydra
from omegaconf import DictConfig
from simwam.runtime_grpo import run_grpo_training

@hydra.main(config_path="../configs", config_name="train_grpo", version_base="1.3")
def main(cfg: DictConfig):
    run_grpo_training(cfg)

if __name__ == "__main__":
    main()

# src/simwam/runtime_grpo.py
def run_grpo_training(cfg: DictConfig):
    # 1. 日志和输出目录
    misc.register_work_dir(cfg.output_dir)
    setup_logging(...)
    
    # 2. 设备和精度
    device = _resolve_train_device()
    model_dtype = _mixed_precision_to_model_dtype(cfg.mixed_precision)
    
    # 3. 构建模型（通过 Hydra instantiate）
    checkpoint_path = cfg.model.get("checkpoint_path", None)
    model_container = OmegaConf.to_container(cfg.model, resolve=True)
    model_container.pop("checkpoint_path", None)
    model = instantiate(model_container, model_dtype=model_dtype, device=device)
    
    # 4. 加载 SFT 权重（warm-start）
    if checkpoint_path:
        ckpt = Path(checkpoint_path)
        model.load_checkpoint(str(ckpt), optimizer=None)
    
    # 5. 配置 GRPO（注入 LoRA、配置采样器）
    model.configure_grpo(cfg.grpo)
    
    # 6. 构建数据集和奖励
    train_ds = instantiate(cfg.data.train)
    reward = instantiate_reward(cfg.grpo)  # NavSimPDMReward
    
    # 7. 构建 GRPO 训练器并开始训练
    trainer = SimWAMGRPOTrainer(model=model, train_dataset=train_ds, reward=reward, cfg=cfg)
    trainer.train()

# src/simwam/runtime_grpo.py:24-87
def create_simwam_grpo(
    model_id, tokenizer_model_id, video_dit_config,
    tokenizer_max_len, load_text_encoder, proprio_dim,
    action_dit_config, action_dit_pretrained_path,
    skip_dit_load_from_pretrain, video_scheduler,
    action_scheduler, loss, mot_checkpoint_mixed_attn,
    redirect_common_files, model_dtype, device,
):
    """构建 SimWAMGRPO（和 runtime.create_simwam 参数一样）。"""
    from .models.wan22.simwam_grpo import SimWAMGRPO
    
    # 把 DictConfig 转成 dict
    video_dit_config = _as_dict(video_dit_config, "video_dit_config", required=True)
    action_dit_config = _as_dict(action_dit_config, "action_dit_config")
    action_scheduler = _as_dict(action_scheduler, "action_scheduler", required=True)
    
    # 校验 action_scheduler 必需的 key
    required_keys = {"train_shift", "infer_shift", "num_train_timesteps"}
    missing = required_keys - set(action_scheduler.keys())
    if missing:
        raise ValueError(f"`action_scheduler` missing keys: {sorted(missing)}.")
    
    # 调用 SimWAMGRPO.from_wan22_pretrained
    return SimWAMGRPO.from_wan22_pretrained(
        device=device, torch_dtype=model_dtype,
        model_id=model_id, tokenizer_model_id=tokenizer_model_id,
        # ... 其余参数
        action_train_shift=float(action_scheduler["train_shift"]),
        action_infer_shift=float(action_scheduler["infer_shift"]),
        action_num_train_timesteps=int(action_scheduler["num_train_timesteps"]),
        loss_lambda_video=float(loss.get("lambda_video", 1.0)),
        loss_lambda_action=float(loss.get("lambda_action", 1.0)),
    )

# src/simwam/models/wan22/action_dit.py
class ActionDiT(nn.Module):
    """
    轻量 Diffusion Transformer，从噪声轨迹生成干净轨迹。
    条件输入：VLM 的缓存 KV + 时间步 + 导航指令 + 自车状态。
    """
    def __init__(self, action_dim=3, hidden_dim=1024, ffn_dim=4096,
                 num_heads=24, attn_head_dim=128, num_layers=30,
                 text_dim=4096, freq_dim=256, eps=1e-6):
        super().__init__()
        
        # 1. 轨迹嵌入：把 (x, y, heading) 的 noisy 轨迹 + Fourier 特征 → hidden_dim
        # 输入维度 = action_dim × 2 + freq_dim = 3×2 + 256 = 262
        self.traj_embed = nn.Sequential(
            nn.Linear(action_dim * 2 + freq_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )
        
        # 2. 时间步嵌入：Flow time → Fourier → MLP → hidden_dim
        self.time_embed = TimestepEmbedding(freq_dim, hidden_dim)
        
        # 3. 文本嵌入：导航指令（T5 输出）→ project → hidden_dim
        self.text_proj = nn.Linear(text_dim, hidden_dim)
        
        # 4. 自车状态嵌入：当前速度/加速度等 → project → hidden_dim
        self.ego_proj = nn.Linear(proprio_dim, hidden_dim)
        
        # 5. 30 层 Transformer Block
        self.blocks = nn.ModuleList([
            ActionDiTBlock(hidden_dim, ffn_dim, num_heads, attn_head_dim, ...)
            for _ in range(num_layers)
        ])
        
        # 6. 输出头：hidden_dim → (action_dim × 2)  # 预测速度场
        self.out_proj = nn.Linear(hidden_dim, action_dim * 2)

# src/simwam/models/wan22/action_dit.py
class ActionDiTBlock(nn.Module):
    def forward(self, x, t_emb, kv_cache, mask):
        # 1. Self-Attention（动作 token 之间互相看）
        #    - 动作 token 是 50 个路点（5s @10Hz），每个 3 维 (x, y, heading)
        #    - 加入因果掩码：后面的路点可以看到前面的，但前面的看不到后面的
        x = x + self.self_attn(self.norm1(x), t_emb, mask=self_causal_mask)
        
        # 2. Cross-Attention（动作 token 看 VLM 的缓存 KV）
        #    - 这是"世界模型信息流入动作专家"的唯一通道
        #    - VLM 的 KV 是从当前观测帧编码的，包含了场景理解
        x = x + self.cross_attn(self.norm2(x), kv_cache, mask=cross_mask)
        
        # 3. FFN（前馈网络）
        x = x + self.ffn(self.norm3(x))
        
        return x

# 轨迹嵌入的输入维度 = action_dim × 2 + freq_dim = 3×2 + 256 = 262
# 其中：
#   - action_dim = 3: (x, y, heading)
#   - × 2: 原始值 + Fourier 编码
#   - freq_dim = 256: Fourier 特征维度

# Fourier 编码的数学：
# pos_enc(x) = [sin(2π·f₁·x), cos(2π·f₁·x), ..., sin(2π·f₂₅₆·x), cos(2π·f₂₅₆·x)]
# 其中 fᵢ = 1/10000^(i/128) 是频率

# src/simwam/models/wan22/mot.py
def build_isolated_mask(video_len, action_len, obs_len):
    """
    构建 isolated attention mask：
    - 观测 tokens：所有 token 都能看到
    - 未来帧 tokens：只有 Video DiT 能看到
    - 动作 tokens：只有 Action DiT 能看到
    """
    total_len = obs_len + video_len + action_len
    
    # 初始化全 1 矩阵（允许所有注意力）
    mask = torch.ones(total_len, total_len, dtype=torch.bool)
    
    # 观测 tokens（前 obs_len 个）：所有 token 都能看到
    mask[:, :obs_len] = True
    
    # 未来帧 tokens（obs_len 到 obs_len+video_len）：
    #   - 只能被自己（Video DiT）看到
    #   - 不能被动作 tokens 看到
    mask[obs_len+video_len:, obs_len:obs_len+video_len] = False  # 动作→未来：禁止
    
    # 动作 tokens（最后 action_len 个）：
    #   - 只能被自己（Action DiT）看到
    #   - 不能被未来帧 tokens 看到
    mask[obs_len:obs_len+video_len, obs_len+video_len:] = False  # 未来→动作：禁止
    
    return mask

# 推理时：直接扔掉 Video DiT 部分，只跑 Action DiT
# 因为动作 tokens 从来不依赖未来帧 tokens
# 所以不需要生成未来帧，大幅减少计算量

# 训练时：两个 DiT 都跑，联合优化
# 推理时：只跑 Action DiT，直接输出轨迹
# → 这就是"训练时借力、推理时独立"

# src/simwam/models/wan22/simwam.py
class SimWAM(nn.Module):
    @classmethod
    def from_wan22_pretrained(cls, model_id, ...):
        """
        从 Wan2.2 预训练权重初始化 SimWAM。
        
        步骤：
        1. 加载 Wan2.2 Video DiT 权重（视频生成 backbone）
        2. 加载 Wan2.2 Video VAE（把图像压缩到 latent）
        3. 加载 T5 文本编码器（编码导航指令）
        4. 初始化 ActionDiT（从预训练权重或随机初始化）
        5. 构建 MOT 联合注意力层
        """
        # 1. Video DiT
        video_dit = WanVideoDiT.from_pretrained(model_id)
        
        # 2. Video VAE
        video_vae = WanVideoVAE.from_pretrained(model_id)
        
        # 3. T5 文本编码器
        text_encoder = WanVideoTextEncoder.from_pretrained(tokenizer_model_id)
        
        # 4. Action DiT
        action_dit = ActionDiT(**action_dit_config)
        if action_dit_pretrained_path:
            action_dit.load_state_dict(torch.load(action_dit_pretrained_path))
        
        # 5. MOT（共享注意力接口）
        mot = MOT(video_dit, action_dit, ...)
        
        return cls(video_dit, action_dit, video_vae, text_encoder, mot, ...)

# src/simwam/models/wan22/simwam.py
def forward(self, images, trajectories, text_input, ...):
    """
    训练时的完整前向传播。
    
    输入：
      - images: 当前帧 + 未来帧的图像 (B, T, C, H, W)
      - trajectories: 专家轨迹 (B, 50, 3)  # 50 路点 × (x, y, heading)
      - text_input: 导航指令的 token IDs
    
    输出：
      - video_loss: 未来帧预测损失
      - action_loss: 轨迹生成损失
    """
    # 1. VAE 编码：图像 → latent
    video_latents = self.video_vae.encode(images)  # (B, T, C', H', W')
    
    # 2. 文本编码：导航指令 → embedding
    text_emb = self.text_encoder(text_input)  # (B, seq_len, 4096)
    
    # 3. 采样时间步 t（Flow Matching）
    t = torch.rand(B, device=devices[0])
    
    # 4. 构造噪声轨迹（Flow Matching 插值）
    noise = torch.randn_like(trajectories)
    noisy_traj = (1 - t) * noise + t * trajectories  # 线性插值
    
    # 5. 联合前向传播（MOT + isolated mask）
    video_pred, action_pred = self.mot(
        video_latents=video_latents,
        noisy_traj=noisy_traj,
        text_emb=text_emb,
        t=t,
        mask=isolated_mask,  # ⭐ 关键：隔离掩码
    )
    
    # 6. 计算损失
    video_loss = F.mse_loss(video_pred, video_latents)
    action_loss = F.mse_loss(action_pred, trajectories - noise)  # 速度场
    
    total_loss = self.lambda_video * video_loss + self.lambda_action * action_loss
    return total_loss, video_loss, action_loss

# src/simwam/models/wan22/simwam.py
@torch.no_grad()
def generate_trajectory(self, images, text_input, num_steps=10):
    """
    推理时：只用 Action DiT，不需要 Video DiT。
    
    流程：
    1. VAE 编码当前帧
    2. T5 编码导航指令
    3. 从高斯噪声初始化轨迹
    4. 10 步欧拉积分 → 干净轨迹
    """
    # 1. 编码当前帧
    obs_latent = self.video_vae.encode(images[:, :1])  # 只编码当前帧
    
    # 2. 编码导航指令
    text_emb = self.text_encoder(text_input)
    
    # 3. 初始化噪声轨迹
    traj = torch.randn(B, 50, 3, device=device)  # 50 路点 × 3
    
    # 4. 10 步欧拉积分
    for i in range(num_steps):
        t = i / num_steps
        # Action DiT 预测速度场
        v_pred = self.action_dit(traj, t, text_emb, obs_latent)
        # 欧拉步
        traj = traj + v_pred * (1.0 / num_steps)
    
    return traj  # (B, 50, 3)  # 最终轨迹

# src/simwam/trainer_grpo.py
class SimWAMGRPOTrainer:
    def train(self):
        for step in range(self.max_steps):
            # ===== Phase 1: 采样（Rollout）=====
            # 从当前策略采样 G=8 条轨迹
            rollout_batch = self.sample_rollouts()
            
            # 计算每条轨迹的奖励
            rewards = self.reward(rollout_batch)  # NAVSIM PDM 分数
            
            # 计算组内归一化 advantage
            advantages = self.compute_advantages(rewards)  
            # advantage_i = (r_i - mean) / (std + eps)
            
            # ===== Phase 2: 训练（Update）=====
            # 复用每批 rollout 做 num_inner_epochs 次更新
            for inner_epoch in range(self.num_inner_epochs):
                self.update_policy(rollout_batch, advantages)

# src/simwam/trainer_grpo.py
def sample_rollouts(self):
    """
    FlowGRPO 的 SDE 采样：在确定性 ODE 上加随机扰动，
    让模型探索不同的轨迹方向。
    
    关键：只在最后 3 步加扰动（k ∈ {7, 8, 9}），
    因为靠近输出端的扰动对最终轨迹影响最直接。
    """
    # 1. 初始化噪声轨迹
    traj = torch.randn(B, 50, 3)
    
    # 2. 10 步欧拉积分，最后 3 步加 SDE 扰动
    for k in range(10):
        t = k / 10
        
        # Action DiT 预测速度场
        v_pred = self.model.action_dit(traj, t, ...)
        
        # ODE 步
        traj = traj + v_pred * (1.0 / 10)
        
        # 如果是最后 3 步（k ∈ {7, 8, 9}），加 SDE 扰动
        if k in [7, 8, 9]:
            # 低频余弦基扰动（不是独立路点噪声）
            noise = self.compute_low_freq_noise(traj)  # 6 个余弦基
            score_correction = self.compute_score_correction(traj, t)
            traj = traj + sigma * noise + 0.5 * sigma**2 * score_correction
    
    return traj

# src/simwam/datasets/navsim/pdm_reward.py
class NavSimPDMReward:
    def __call__(self, trajectories):
        """
        计算 NAVSIM PDM 分数。
        
        PDMS = NC × DAC × (5·TTC + 5·EP + 2·C) / 12
        
        NC: 无责碰撞（0/0.5/1）
        DAC: 可行驶区域（0/1）
        TTC: 碰撞时间（0/1）
        EP: 前进进度（0-1 连续）
        C: 舒适度（0/1）
        """
        rewards = []
        for traj in trajectories:
            # 把轨迹转换成 NAVSIM 格式
            navsim_traj = self.convert_to_navsim(traj)
            
            # 调用 NAVSIM 评测器
            pdms = self.pdm_scorer(navsim_traj)
            
            rewards.append(pdms)
        
        return torch.tensor(rewards)

# src/simwam/trainer_grpo.py
def compute_advantages(self, rewards):
    """
    组内归一化：把奖励翻译成 advantage。
    
    A_i = (r_i - μ) / σ
    
    其中 μ = mean(rewards), σ = std(rewards) + 1e-4
    """
    mean = rewards.mean()
    std = rewards.std() + 1e-4
    advantages = (rewards - mean) / std
    
    # 截断到 ±adv_clip_max
    advantages = advantages.clamp(-self.adv_clip_max, self.adv_clip_max)
    
    return advantages

# src/simwam/trainer_grpo.py
def update_policy(self, rollout_batch, advantages):
    """
    PPO 风格的策略更新。
    
    核心：用 ratio = π_new / π_old 控制更新幅度，
    加 clip 防止一步更新太猛。
    """
    # 计算新模型下的 log_prob
    log_probs_new = self.model.compute_log_prob(rollout_batch)
    
    # PPO ratio
    ratio = torch.exp(log_probs_new - rollout_batch.log_probs_old)
    
    # PPO clipped loss
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1 - self.ppo_clip_range, 1 + self.ppo_clip_range) * advantages
    
    loss = -torch.min(surr1, surr2).mean()
    
    # 反向传播（只更新 LoRA 参数）
    loss.backward()
    self.optimizer.step()
    self.optimizer.zero_grad()

# src/simwam/models/wan22/lora.py
class LoRALinear(nn.Module):
    """
    低秩适配器：只训练 A 和 B 两个小矩阵，
    冻结原始权重 W。
    
    W' = W + (α/r) · B @ A
    """
    def __init__(self, in_features, out_features, r=16, alpha=32.0):
        self.lora_A = nn.Linear(in_features, r, bias=False)
        self.lora_B = nn.Linear(r, out_features, bias=False)
        self.alpha = alpha
        self.scaling = alpha / r  # = 32/16 = 2
    
    def forward(self, x):
        # 原始输出 + LoRA 增量
        return F.linear(x, self.weight) + self.scaling * self.lora_B(self.lora_A(x))

# experiments/navsim/run_eval_navsim.sh
CKPT=./runs/grpo/.../checkpoints/weights/step_2000.pt \
TASK=navsim_grpo_action_pdm_384x672_flowgrpo_lora \
NPROC_PER_NODE=8 \
bash experiments/navsim/run_eval_navsim.sh

# experiments/navsim/eval_navsim.py
def evaluate(checkpoint_path, task_config):
    # 1. 加载模型
    model = load_model(checkpoint_path)
    
    # 2. 加载 NAVSIM navtest 数据集
    test_dataset = NavSimDataset(split="navtest")
    
    # 3. 对每个场景生成轨迹
    for scene in test_dataset:
        traj = model.generate_trajectory(scene.images, scene.navigation)
        
        # 4. 用 NAVSIM PDM 评测器打分
        pdms = pdm_scorer(traj, scene)
        
        # 5. 保存结果
        results.append({
            "scene": scene.token,
            "trajectory": traj,
            "nc": pdms.nc,
            "dac": pdms.dac,
            "ep": pdms.ep,
            "ttc": pdms.ttc,
            "comfort": pdms.comfort,
            "pdms": pdms.score,
        })
    
    # 6. 汇总统计
    avg_pdms = np.mean([r["pdms"] for r in results])
    print(f"PDMS: {avg_pdms:.1f}")