looyifan / HiP-AD: Hierarchical and Multi-granularity Planning with Deformable Attention

import torch
import torch.nn.functional as F

def Unified_Decoder(
    image_features,   # [多视角环视图像特征 F] (监控录像)
    queries,          # [动态物体Qa, 静态地图Qm, 规划主脑Qp] (拿着清单的代表)
    anchors,          # [动态框Aa, 静态线段Am, 未来轨迹Ap] (各自关注的3D物理坐标)
    history_memory,   # [上一帧存下来的重要记忆]
    camera_params,    # [相机的内外参矩阵] (用于3D到2D的投影)
    num_layers=6      # 会议通常要开好几轮 (堆叠6层Decoder)
):
    # --- 会议前奏：代表入场 ---
    Qa, Qm, Qp = queries
    Aa, Am, Ap = anchors

    # 🔄 开始循环开会：每一层 Decoder 都在进行信息的深度融合
    for layer in range(num_layers):

        # =========================================================
        # 🕒 议程一：时序交互模块 (Temporal Interaction Module)
        # 目标："温故知新"，看看上一秒发生了什么。
        # =========================================================

        # 1. 脑容量有限，用 Top-k 机制挑选上一帧最有价值的记忆点
        hist_k_agent, hist_k_map, hist_k_plan = select_top_k(history_memory)

        # 2. 感知代表各自查阅自己的历史卷宗 (Cross-Attention)
        Qa = CrossAttention(query=Qa, key=hist_k_agent, value=hist_k_agent)
        Qm = CrossAttention(query=Qm, key=hist_k_map,   value=hist_k_map)

        # 💡 3. 【老司机的特权】规划主脑不仅看自己的历史轨迹，
        #    还要盯着历史的感知环境！
        Qp_self_hist = CrossAttention(query=Qp, key=hist_k_plan, value=hist_k_plan)
        Qp_perc_hist = CrossAttention(
            query=Qp,
            key=concat(hist_k_agent, hist_k_map),
            value=concat(hist_k_agent, hist_k_map)
        )
        Qp = combine_features(Qp_self_hist, Qp_perc_hist)

        # =========================================================
        # 🗣️ 议程二：协同交互模块 (Collaborative Interaction Module)
        # 目标：统一的圆桌会议，互相通气，引入"物理距离惩罚"！
        # =========================================================

        # 1. 把所有人拉进一个群里
        Q_all = concat(Qa, Qm, Qp)
        A_all = concat(Aa, Am, Ap)

        # 2. 计算纯粹的注意力分数 (谁跟谁有关联)
        attn_scores = torch.matmul(Q_all, Q_all.transpose(-2, -1)) / math.sqrt(C)

        # 💡 3. 【核心技术细节：公式(1)】计算真实物理世界中 3D 锚点之间的距离矩阵 D
        D_matrix = calc_euclidean_distance_3d(A_all, A_all)

        # ⚠️ 【主脑特权掩码】规划主脑(Qp)看所有人都不受距离限制
        D_matrix = apply_planning_mask(D_matrix, mask_value=0.0)

        # 4. 用一个小网络学习一个动态系数 tau (控制惩罚力度)
        tau = MLP(Q_all)

        # 5. 施加物理距离惩罚！距离越远的物体，注意力得分被扣得越惨
        penalized_scores = attn_scores - (tau * D_matrix)
        attn_weights = F.softmax(penalized_scores, dim=-1)

        # 6. 根据惩罚后的权重，大家交换情报
        Q_all_updated = torch.matmul(attn_weights, Q_all)

        # 7. 散会，大家拿着更新后的情报各自归位
        Qa, Qm, Qp = split_queries(Q_all_updated)

        # =========================================================
        # 🎯 议程三：任务可变形注意力模块 (Task Deformable Attention Module)
        # 目标：拿着讨论结果，去图像画面里"精准取证"。
        # =========================================================

        # 1. 感知代表去找车和线：把 3D 的框和线投影到 2D 图像上
        P_a_2d = project_3D_to_2D(Aa, camera_params)
        P_m_2d = project_3D_to_2D(Am, camera_params)
        Qa = DeformableAttention(query=Qa, reference_points=P_a_2d, features=image_features)
        Qm = DeformableAttention(query=Qm, reference_points=P_m_2d, features=image_features)

        # 💡 2. 【核心技术细节：公式(2) PDA】规划主脑的终极绝技！
        #    把预测的"未来3D行驶轨迹路点"投影到多个摄像头画面上
        P_p_2d = project_trajectory_to_2D(Ap, camera_params, predefined_heights)

        # 让主脑自己学习：在轨迹周围哪些像素点最值得关注？
        sampling_offsets, sampling_weights = MLP_predict_offsets_weights(Qp)

        # 施展 PDA：目光死死盯住"车轮即将压过的未来路线"周围的画面！
        Qp = PDA_DeformableAttention(
            query=Qp,
            reference_points=P_p_2d,
            offsets=sampling_offsets,
            weights=sampling_weights,
            features=image_features
        )

        # 🔄 层级收尾：更新代表们的 3D 坐标锚点
        Aa = update_anchors(Aa, Qa)
        Am = update_anchors(Am, Qm)
        Ap = update_anchors(Ap, Qp)  # 轨迹越来越精确

    # 🚪 整个开会流程结束，输出给外面的业务部门 (Heads) 去执行
    update_history_memory(Qa, Qm, Qp)

    return Qa, Qm, Qp