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):
    """完整 SDE 采样流程（带详细注释）。"""
    
    # ========== 第一步：初始化噪声 ==========
    # 从标准正态分布采样初始轨迹（50个路点 × 3维：x, y, heading）
    # 这就是 Flow Matching 的起点 x_0 ~ N(0, I)
    traj = torch.randn(B, 50, 3)  # (batch=8, waypoints=50, dim=3)
    
    # ========== 第二步：10步欧拉积分 ==========
    for k in range(10):
        t = k / 10  # 时间步 t ∈ {0.0, 0.1, ..., 0.9}
        
        # --- 纯 ODE 步 ---
        # Action DiT 预测当前时刻的速度场 v_θ(x_k, t, c)
        # 输入：带噪轨迹 + 时间步 + VLM 条件（cross-attention KV）
        v_pred = self.model.action_dit(traj, t, **conditions)  # (B, 50, 3)
        
        # 确定性 ODE 步：x_{k+1} = x_k + v_θ · Δt
        traj_ode = traj + v_pred * (1.0 / 10)  # Δt = 1/10 = 0.1
        
        # --- SDE 步（仅 k ∈ {7, 8, 9}）---
        if k in [7, 8, 9]:
            # 计算噪声级别 σ（可能随时间步变化）
            sigma = self.compute_sigma(t)  # 通常 σ = noise_level = 0.1
            
            # 计算低频余弦基噪声（不是独立路点噪声！）
            noise = self.compute_low_freq_noise(traj)  # (B, 50, 3)
            
            # 计算 score correction（概率梯度修正）
            score_correction = self.compute_score_correction(traj, t, v_pred)
            # score_correction ≈ -[x_k + (1-t)·v_θ] / t
            
            # SDE 更新：ODE步 + 噪声 + score修正
            traj = traj_ode + sigma * noise + 0.5 * sigma**2 * score_correction
        else:
            # 非SDE步：直接用ODE结果
            traj = traj_ode
    
    return traj  # (B, 50, 3) → 8条采样轨迹

# ❌ 独立噪声：每个路点独立扰动
noise_independent = torch.randn(50, 3)  # 50个路点各自独立
# 结果：第3个路点突然左拐0.5m，第4个又右拐0.3m → 高频抖动
# 物理上不可执行：车辆的转向系统无法做这种突变

# ✅ 低频余弦基：用6个平滑基函数组合
def compute_low_freq_noise(traj):
    """用余弦基函数生成平滑噪声。"""
    # 路点索引 i = 0, 1, ..., 49
    i = torch.arange(50).float()
    
    # 6个基函数：cos(π·i/50 · k), k=1,2,...,6
    # 每个基函数都是全局平滑的
    bases = torch.stack([torch.cos(torch.pi * i / 50 * k) for k in range(1, 7)])
    # bases shape: (6, 50)
    
    # 随机权重
    weights = torch.randn(6)  # 6个权重
    
    # 组合：noise(i) = Σ w_k · cos(π·i·k/50)
    noise_1d = weights @ bases  # (50,)
    
    # 扩展到3维（x, y, heading 分别用不同权重）
    noise = torch.stack([noise_1d] * 3, dim=-1)  # (50, 3)
    
    return noise

def compute_nc(ego_traj, agent_trajs, agent_metadata):
    """
    NC = 0.0  如果自车是碰撞的主责方
    NC = 0.5  如果自车是碰撞的次责方
    NC = 1.0  如果没有碰撞，或碰撞完全由对方造成
    
    判断依据：碰撞点的位置 + 双方的速度方向
    """
    for agent in agent_trajs:
        # 检测 OBB（有向包围盒）重叠
        collision_point = detect_obb_overlap(ego_traj, agent)
        
        if collision_point is not None:
            # 判断责任：谁的运动方向指向碰撞点？
            ego_heading_to_collision = angle(ego_vel, collision_point - ego_pos)
            agent_heading_to_collision = angle(agent_vel, collision_point - agent_pos)
            
            if ego_heading_to_collision < agent_heading_to_collision:
                return 0.0  # 自车主责
            else:
                return 0.5  # 自车次责
    
    return 1.0  # 无碰撞

def compute_dac(ego_traj, road_boundaries):
    """
    DAC = 1.0  所有路点都在可行驶区域内
    DAC = 0.0  任何一个路点越界
    
    判断方法：每个路点 vs 道路边界的多边形包含测试
    """
    for waypoint in ego_traj:
        if not is_inside_drivable_area(waypoint, road_boundaries):
            return 0.0
    return 1.0

def compute_ttc(ego_traj, agent_trajs):
    """
    TTC ∈ [0, 1]
    
    TTC = 1.0  TTC > 5s（非常安全）
    TTC = 0.0  TTC < 0.5s（即将碰撞）
    TTC = 中间值  线性插值
    
    计算方法：对每对(自车, 障碍物)，找最近的碰撞时间
    """
    min_ttc = float('inf')
    
    for agent in agent_trajs:
        # 简化模型：假设双方匀速运动，求最近距离时刻
        rel_pos = ego_pos - agent.pos
        rel_vel = ego_vel - agent.vel
        
        # 最近时刻 t* = -rel_pos · rel_vel / |rel_vel|²
        t_star = -dot(rel_pos, rel_vel) / (dot(rel_vel, rel_vel) + 1e-6)
        t_star = max(0, t_star)
        
        # 最近距离
        min_dist = norm(rel_pos + rel_vel * t_star)
        
        # 碰撞时间：最近距离 < 安全阈值时的最早时刻
        if min_dist < safety_threshold:
            ttc = t_star
            min_ttc = min(min_ttc, ttc)
    
    # 归一化到 [0, 1]
    if min_ttc < 0.5:
        return 0.0
    elif min_ttc > 5.0:
        return 1.0
    else:
        return (min_ttc - 0.5) / 4.5  # 线性插值

def compute_ep(ego_traj, max_progress=60.0):
    """
    EP ∈ [0, 1]
    
    EP = 自车在规划窗口内前进的距离 / 最大期望前进距离
    
    例：规划 5s，自车前进了 30m，最大期望 60m → EP = 0.5
    """
    # 计算自车在规划窗口内的纵向位移
    progress = ego_traj[-1].y - ego_traj[0].y  # y轴 = 前进方向
    
    # 归一化
    ep = min(progress / max_progress, 1.0)
    return max(ep, 0.0)

def compute_comfort(ego_traj, dt=0.1):
    """
    C = 1.0  横向加速度 < 0.5m/s² 且 纵向加速度 < 1.5m/s²
    C = 0.0  任何一个超阈值
    
    实际是二值的：舒适 or 不舒适
    """
    for i in range(1, len(ego_traj)):
        # 横向加速度（转向离心力）
        lat_acc = compute_lateral_acceleration(ego_traj, i, dt)
        # 纵向加速度（急加速/急刹车）
        lon_acc = compute_longitudinal_acceleration(ego_traj, i, dt)
        
        if abs(lat_acc) > 0.5 or abs(lon_acc) > 1.5:
            return 0.0
    return 1.0

# src/simwam/datasets/navsim/pdm_reward.py
class NavSimPDMReward:
    def __call__(self, trajectories):
        """对每条采样轨迹打 PDM 分数。"""
        rewards = []
        for traj in trajectories:
            # 1. 坐标转换：SimWAM 格式 → NAVSIM 格式
            navsim_traj = self.convert_to_navsim(traj)
            
            # 2. 逐项计算
            nc = self.compute_nc(navsim_traj)
            dac = self.compute_dac(navsim_traj)
            ttc = self.compute_ttc(navsim_traj)
            ep = self.compute_ep(navsim_traj)
            c = self.compute_comfort(navsim_traj)
            
            # 3. 加权组合
            pdms = nc * dac * (5*ttc + 5*ep + 2*c) / 12.0
            
            rewards.append({
                "pdms": pdms,
                "nc": nc, "dac": dac, "ttc": ttc, "ep": ep, "c": c
            })
        
        return torch.tensor([r["pdms"] for r in rewards])

# 8条轨迹的PDM奖励
rewards = torch.tensor([0.85, 0.72, 0.91, 0.68, 0.83, 0.77, 0.88, 0.79])

# 第一步：算均值和标准差
mean = rewards.mean()  # = 0.80375
std = rewards.std() + 1e-4  # = 0.08214

# 第二步：组内归一化
advantages = (rewards - mean) / std
# advantages = [ 0.56, -1.02,  1.29, -1.48,  0.32, -0.41,  0.93, -0.20]

adv_clip_max = 5.0
advantages = advantages.clamp(-adv_clip_max, adv_clip_max)
# 上面的例子没有超限，但如果有极端奖励（比如 PDM=100），
# 截断防止梯度爆炸

# 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

# PPO clipped loss 的直觉
# 对于好轨迹（A>0）：
#   ratio < 1-ε: 被 clip → 阻止"过度降低"好轨迹的概率
#   1-ε < ratio < 1+ε: 正常更新
#   ratio > 1+ε: 被 clip → 阻止"过度提高"好轨迹的概率（防止reward hacking）

# 对于差轨迹（A<0）：
#   ratio < 1-ε: 被 clip → 阻止"过度降低"差轨迹的概率
#   1-ε < ratio < 1+ε: 正常更新
#   ratio > 1+ε: 被 clip → 阻止"过度提高"差轨迹的概率

# src/simwam/trainer_grpo.py
def update_policy(self, rollout_batch, advantages):
    """
    PPO 风格的策略更新。
    """
    # 1. 用新模型重新计算 log_prob
    # rollout_batch 中保存了采样时的 log_probs_old
    log_probs_new = self.model.compute_log_prob(rollout_batch)
    # log_probs_new shape: (B, T) → 每条轨迹每步的 log_prob
    
    # 2. 计算轨迹级 log_prob（所有步求和）
    log_prob_new = log_probs_new.sum(dim=-1)   # (B,)
    log_prob_old = rollout_batch.log_probs_old.sum(dim=-1)  # (B,)
    
    # 3. 计算 ratio
    ratio = torch.exp(log_prob_new - log_prob_old)
    # ratio = π_new(τ) / π_old(τ)
    # ratio > 1 → 新策略更偏好这条轨迹
    # ratio < 1 → 新策略不太偏好这条轨迹
    
    # 4. PPO clipped surrogate 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()
    # 直觉：如果 clip 没起作用（ratio 在范围内），loss = -ratio·A
    #       如果 clip 起作用（ratio 超范围），loss = -clip(ratio)·A（更保守）
    
    # 5. 反向传播（只更新 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}")

# flow_grpo/rewards.py
reward_fn = multi_score(device, config.reward_fn)
# 返回 {"avg": [0.45, 0.78, ...], "strict_accuracy": [0, 1, ...]}

# 具体奖励类型：
# - OCR reward: 图中文字是否正确渲染
# - PickScore: 人类偏好分数
# - GenEval: 组合逻辑检查（几个物体、什么颜色、什么关系）

# src/simwam/datasets/navsim/pdm_reward.py
PDMS = NC × DAC × (5·TTC + 5·EP + 2·C) / 12

# flux_pipeline_with_logprob_fast.py
# 窗口在 [sde_window_range[0], sde_window_range[1]] 内随机选取
start = randint(sde_window_range[0], sde_window_range[1] - sde_window_size)
end = start + sde_window_size
sde_window = (start, end)  # 例如 (2, 5)

# 窗口前：纯 ODE（无噪声）
# 窗口内：SDE（加噪声）
# 窗口后：纯 ODE

# src/simwam/trainer_grpo.py
for k in range(10):  # 10 步欧拉积分
    if k in [7, 8, 9]:  # 只在最后 3 步加 SDE 扰动
        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

# 标准 SDE：在 latent 空间加各向同性高斯噪声
variance_noise = randn_tensor(latents.shape)
prev_sample = prev_sample_mean + std_dev_t * sqrt(-dt) * variance_noise

# 不是独立路点噪声，而是限制在 6 个余弦基上
# 只在"整体偏左/偏右""整体加速/减速"这几个低维模态上探索
noise = compute_low_freq_noise(traj)  # 6 个余弦基

# 完整 loss = PPO policy loss + KL penalty
policy_loss = torch.mean(torch.maximum(unclipped_loss, clipped_loss))

# KL penalty：当前策略 vs 参考策略（关闭 LoRA 后的基座模型）
with transformer.module.disable_adapter():  # 关闭 LoRA = 参考模型
    _, _, prev_sample_mean_ref, _ = compute_log_prob(...)

kl_loss = ((prev_sample_mean - prev_sample_mean_ref) ** 2).mean() / (2 * std_dev_t ** 2)
loss = policy_loss + config.train.beta * kl_loss

# SimWAM 只用 PPO clipped loss，没有 KL 项
ratio = torch.exp(log_probs_new - rollout_batch.log_probs_old)
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1 - self.ppo_clip_range, 1 + self.ppo_clip_range) * advantages
loss = -torch.min(surr1, surr2).mean()

# FLUX 12B 模型，需要更大容量的 LoRA
transformer_lora_config = LoraConfig(
    r=64,              # rank=64
    lora_alpha=128,    # alpha=128，scaling=2
    target_modules=["attn.to_k", "attn.to_q", "attn.to_v", "attn.to_out.0", ...]
)

# ActionDiT 参数量小，LoRA 也小
lora:
  r: 16              # rank=16（FLUX 的 1/4）
  alpha: 32.0        # alpha=32，scaling=2（比例相同）
  target_modules: [q, k, v, o]  # 只在注意力层

# 每 8 步更新 EMA 参数
ema_decay = 0.9
ema_params = decay * ema_params + (1 - decay) * current_params

# 训练用实时参数，评估/保存用 EMA 参数