src/wjad/model.py

"""端到端自动驾驶模型 E2EAVModel。

forward 流程
  1. ``DINOv3`` 提取 8 帧 patch 特征。
  2. ``OnlineCalibration`` 用原始 ego/intr/extr (symlog) + DINOv3 patch 作 KV，
     输出 symlog 空间残差，叠加并 symexp 还原得到 corrected_*。
  3. 用 corrected_intr / corrected_extr / corrected_ego 计算
     - 每 token 的自车系单位射线（仅用于视觉 token 的 RoPE 第一组头）。
     - 8 个 ego token（symlog 后线性投影）。
  4. 2×2×2 时空压缩 -> 1536 视觉 token。
  5. 拼接 [vision(1536) | ego(8) | det(1024) | ctrl(24) | extra(256)] = 2848 token。
     非视觉切片各自加可学习 PE。
  6. 18 层主干（仅视觉切片应用 3D RoPE）。
  7. 切片送入 ``DetectionTrajHead`` 与 ``ControlHead``。
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from .backbone import Backbone, BackboneOutput
from .calibration import OnlineCalibration, CalibrationOutput
from .encoders import DINOv3Wrapper
from .heads import (
    ControlHead,
    ControlOutput,
    DetectionTrajHead,
    DetectionTrajOutput,
)
from .modules.learned_pe import LearnedTokenPE
from .modules.normalization import symlog
from .modules.pos_encoding import RoPE3D
from .modules.rays import compute_ego_rays
from .modules.temporal_compress import TemporalCompress2x2x2


@dataclass
class E2EOutput:
    """模型完整输出。"""

    detection: DetectionTrajOutput
    control: ControlOutput
    backbone_out: BackboneOutput
    calibration: CalibrationOutput


class E2EAVModel(nn.Module):
    def __init__(
        self,
        dinov3_path: str = "./dinov3-vitb16-pretrain-lvd1689m",
        backbone_dim: int = 768,
        num_heads: int = 12,
        num_dense_layers: int = 9,
        num_moe_layers: int = 9,
        num_routed_experts: int = 7,
        num_shared_experts: int = 1,
        topk_experts: int = 3,
        ffn_mult: int = 4,
        # token 数量
        num_history_frames: int = 8,
        num_detection_tokens: int = 1024,
        num_control_tokens: int = 24,
        num_ego_tokens: int = 8,
        num_extra_tokens: int = 256,
        # 输入分辨率
        image_h: int = 384,
        image_w: int = 1024,
        patch_size: int = 16,
        # 头超参
        num_classes: int = 22,
        traj_horizon: int = 24,
        det_head_hidden: int = 384,
        ctrl_head_hidden: int = 384,
        # 校准
        calib_dim: int = 256,
        calib_num_query: int = 256,
        calib_num_blocks: int = 2,
        calib_num_self_per_block: int = 2,
        calib_num_heads: int = 8,
        calib_residual_range: float = 0.1,
        calib_intr_dim: int = 11,
        # DINOv3
        freeze_dinov3: bool = True,
        attn_implementation: str = "sdpa",
    ) -> None:
        super().__init__()
        self.image_h = image_h
        self.image_w = image_w
        self.patch_size = patch_size
        self.num_history = num_history_frames
        self.num_det = num_detection_tokens
        self.num_ctrl = num_control_tokens
        self.num_ego = num_ego_tokens
        self.num_extra = num_extra_tokens

        # === 1) DINOv3 ===
        self.dinov3 = DINOv3Wrapper(
            pretrained_path=dinov3_path,
            attn_implementation=attn_implementation,
            freeze=freeze_dinov3,
        )
        dino_dim = self.dinov3.hidden_size

        # === 2) 在线校准 ===
        self.calib = OnlineCalibration(
            dino_dim=dino_dim,
            hidden_dim=calib_dim,
            num_query_tokens=calib_num_query,
            num_blocks=calib_num_blocks,
            num_self_attn_per_block=calib_num_self_per_block,
            num_heads=calib_num_heads,
            residual_range=calib_residual_range,
            num_history_frames=num_history_frames,
            intr_dim=calib_intr_dim,
        )

        # === 3) 时空压缩 ===
        self.compress = TemporalCompress2x2x2(dim=dino_dim)
        # patch 网格大小（必须能被 2 整除）
        self.gh = image_h // patch_size
        self.gw = image_w // patch_size

        # === 4) 各类 token + 可学习 PE ===
        self.ego_proj = nn.Linear(6, backbone_dim)  # 6D pose -> backbone dim
        self.det_tokens = nn.Parameter(torch.empty(num_detection_tokens, backbone_dim))
        nn.init.trunc_normal_(self.det_tokens, std=0.02)
        self.ctrl_tokens = nn.Parameter(torch.empty(num_control_tokens, backbone_dim))
        nn.init.trunc_normal_(self.ctrl_tokens, std=0.02)
        self.extra_tokens = nn.Parameter(torch.empty(num_extra_tokens, backbone_dim))
        nn.init.trunc_normal_(self.extra_tokens, std=0.02)

        self.ego_pe = LearnedTokenPE(num_ego_tokens, backbone_dim)
        self.det_pe = LearnedTokenPE(num_detection_tokens, backbone_dim)
        self.ctrl_pe = LearnedTokenPE(num_control_tokens, backbone_dim)
        self.extra_pe = LearnedTokenPE(num_extra_tokens, backbone_dim)

        # === 5) RoPE 3D（仅视觉，4 时间帧 × 12 × 32 网格）===
        self.rope = RoPE3D(
            num_heads=num_heads,
            head_dim=backbone_dim // num_heads,
            time_size=num_history_frames // 2,
            height_size=self.gh // 2,
            width_size=self.gw // 2,
        )

        # === 6) 主干 18 层 ===
        self.backbone = Backbone(
            dim=backbone_dim,
            num_heads=num_heads,
            ffn_mult=ffn_mult,
            num_dense_layers=num_dense_layers,
            num_moe_layers=num_moe_layers,
            num_routed=num_routed_experts,
            num_shared=num_shared_experts,
            topk=topk_experts,
        )

        # === 7) 头 ===
        self.det_traj_head = DetectionTrajHead(
            in_dim=backbone_dim,
            hidden_size=det_head_hidden,
            num_classes=num_classes,
            traj_horizon=traj_horizon,
        )
        self.ctrl_head = ControlHead(
            in_dim=backbone_dim,
            hidden_size=ctrl_head_hidden,
            num_traj_tokens=12,
            num_action_tokens=num_control_tokens - 12,
            ego_traj_horizon=traj_horizon,
        )

    # ---------- 工具 ----------

    @property
    def num_visual_tokens(self) -> int:
        # 2×2×2 压缩后
        return (self.num_history // 2) * (self.gh // 2) * (self.gw // 2)

    def _build_ego_tokens(self, ego_6d_corrected: torch.Tensor) -> torch.Tensor:
        """``[B, 8, 6]`` -> symlog -> Linear -> ``[B, 8, D]``。"""
        return self.ego_proj(symlog(ego_6d_corrected))

    def _build_visual_rays(
        self,
        intr_corrected: torch.Tensor,    # [B, calib_intr_dim]
        extr_corrected_se3: torch.Tensor, # [B, 4, 4] cam2vehicle
        compressed_thw: tuple[int, int, int],
    ) -> torch.Tensor:
        """计算压缩后视觉 token 网格的射线方向。

        在 2×2×2 压缩后，每个视觉 token 对应原 patch 网格的一个 2x2 区域 +
        2 个时间帧。这里取所代表区域的中心像素与"中间时间"的射线作近似，
        所有时间帧取同一个 (h, w) 上的射线（因为相机 pose 在 8 帧间是
        rigid 的相机系；自车运动差异会通过 ego token 传递）。
        """
        b = intr_corrected.shape[0]
        t_, h_, w_ = compressed_thw
        rays_grid = compute_ego_rays(
            intr_vec=intr_corrected,
            cam2vehicle=extr_corrected_se3,
            height=self.image_h,
            width=self.image_w,
            grid_h=h_,
            grid_w=w_,
            device=intr_corrected.device,
            dtype=intr_corrected.dtype,
        )  # [B, h_, w_, 3]
        # 复制到时间维：[B, T_, h_, w_, 3] -> flatten 为 [B, N_v, 3]
        rays = rays_grid.unsqueeze(1).expand(-1, t_, -1, -1, -1).contiguous()
        rays = rays.reshape(b, t_ * h_ * w_, 3)
        return rays

    # ---------- 前向 ----------

    def forward(
        self,
        images: torch.Tensor,        # [B, T=8, 3, H, W]
        ego_6d_raw: torch.Tensor,     # [B, 8, 6]
        intr_raw: torch.Tensor,        # [B, calib_intr_dim]，须与构造时一致
        extr_6d_raw: torch.Tensor,     # [B, 6]
    ) -> E2EOutput:
        b, t, _, h, w = images.shape
        assert t == self.num_history, f"history frames mismatch: {t} vs {self.num_history}"

        # 1) DINOv3 patch tokens [B, T, gh, gw, D_dino]
        dino_feats = self.dinov3(images)

        # 2) 校准（symlog 空间残差 + symexp 还原）
        calib_out: CalibrationOutput = self.calib(
            dino_feats=dino_feats,
            ego_raw=ego_6d_raw,
            intr_raw=intr_raw,
            extr_raw=extr_6d_raw,
        )
        corrected_ego = calib_out.corrected_ego
        corrected_intr = calib_out.corrected_intr
        corrected_extr_6d = calib_out.corrected_extr

        # 3) 把 corrected_extr 6D 转成 4x4
        from .data.se3 import six_d_to_matrix
        cam2veh_corrected = six_d_to_matrix(corrected_extr_6d)  # [B, 4, 4]

        # 4) 2x2x2 时空压缩
        compressed, thw = self.compress(dino_feats)  # [B, N_v, D]
        n_v = compressed.shape[1]

        # 5) 视觉射线（用 corrected_intr / corrected_extr）
        rays = self._build_visual_rays(corrected_intr, cam2veh_corrected, thw)
        rope_cos, rope_sin = self.rope.compute_freqs(rays)

        # 6) 构造非视觉 token
        ego_tok = self._build_ego_tokens(corrected_ego)         # [B, 8, D]
        det_tok = self.det_tokens.unsqueeze(0).expand(b, -1, -1)
        ctrl_tok = self.ctrl_tokens.unsqueeze(0).expand(b, -1, -1)
        extra_tok = self.extra_tokens.unsqueeze(0).expand(b, -1, -1)

        ego_tok = self.ego_pe(ego_tok)
        det_tok = self.det_pe(det_tok)
        ctrl_tok = self.ctrl_pe(ctrl_tok)
        extra_tok = self.extra_pe(extra_tok)

        # 7) 拼接序列：[vision | ego | det | ctrl | extra]
        seq = torch.cat([compressed, ego_tok, det_tok, ctrl_tok, extra_tok], dim=1)
        visual_slice = (0, n_v)

        # 8) 主干
        bb_out = self.backbone(seq, rope_cos=rope_cos, rope_sin=rope_sin, visual_slice=visual_slice)

        # 9) 切片送入头
        offset_det = n_v + self.num_ego
        offset_ctrl = offset_det + self.num_det

        det_feats = bb_out.hidden_states[:, offset_det : offset_det + self.num_det]
        ctrl_feats = bb_out.hidden_states[:, offset_ctrl : offset_ctrl + self.num_ctrl]

        det_out = self.det_traj_head(det_feats)
        ctrl_out = self.ctrl_head(ctrl_feats)

        return E2EOutput(
            detection=det_out,
            control=ctrl_out,
            backbone_out=bb_out,
            calibration=calib_out,
        )