4B 干翻 GPT-6 Astra！国产 grounding 基础模型把 34 个感知基准刷成 SOTA，还顺手登顶 HF 日榜第一

# -*- coding: utf-8 -*-
"""
GroundingPI 全流程伪代码
配套文章:《4B 打赢 GPT-6 Astra:小鹏 GroundingPI 把 34 项感知基准刷成 SOTA》

论文: GroundingPI: A Grounding Foundation Model towards Physical Intelligence
      with Visual Primitives (arXiv:2609.39601v2)
开发方: 小鹏集团 (XPeng Inc.) 联合北京大学、香港大学等高校团队

--------------------------------------------------------------------------
这个文件是「可执行的说明文档」,不是训练工程:

  * 建模 / 数据引擎 / 训练 / 评测 / 物理智能 各节用抽象接口 + 注释表达,
    函数体留空 (...),代表「这里应为真实实现」。
  * 只有第 2 节 GAM 空间词元协议是真实可运行的纯 Python,
    用来把全文最容易写错的一处协议讲清楚(相邻坐标词元之间不插空格)。
  * main() 不触发任何训练,只打印流程骨架 + 跑一遍协议自检。

运行: python3 groundingpi_pseudocode.py
--------------------------------------------------------------------------
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any


# ============================================================================
# 0. 全局配置:只写论文 / README 中可核对的数字,不补未披露的量
# ============================================================================

@dataclass(frozen=True)
class Config:
    # --- 模型 ---
    visual_encoder: str = "MoonViT-V2 (Kimi K3)"   # 上游组件,非开发方
    projector: str = "learnable multimodal projector"
    llm_backbone: str = "Qwen3-4B"                  # 上游组件,非开发方
    params: str = "~4B"

    # --- GAM 空间词元协议 ---
    coord_bins: int = 1000        # <0> ~ <999>,与语义标签共享同一输出词表
    coord_max: int = 999
    grid_w: int = 1920            # 演示用图像尺寸,仅用于坐标还原
    grid_h: int = 1080

    # --- 训练:三大阶段 / 四个 Stage ---
    stage1: str = "Projector 对齐(冻结 ViT 与 LLM,image-caption,8192)"
    stage2: str = "联合多模态预训练(全模块,text-only + 图文,8192)"
    stage3: str = "通用视觉与视频理解(全模块,32768)"
    stage4: str = "坐标对齐 / SFT(8 个空间任务组,8192)"
    stage5: str = "GRPO 强化后训练(G=8,eps=0.2,beta=0.02,冻结 SFT 参考,只更新语言参数)"
    grpo_group: int = 8

    # --- 推理服务(README 默认配置) ---
    serving: str = "vLLM 0.18.x / BF16 / TP=1 / 单图 / 16384 上下文 / eager(禁用 CUDA Graph)"
    decode_flags: str = "skip_special_tokens=False, spaces_between_special_tokens=False"

    # --- 评测口径(注意:三个数字口径不同,不要混用) ---
    n_benchmarks: int = 34        # 34 个感知基准
    n_baselines: int = 44         # 44 个对比基线
    n_capabilities: int = 11      # 11 类感知能力(等权平均)
    n_spatial_groups: int = 8     # 训练 Stage 4 的 8 个空间任务组,与上者不同
    avg_score: float = 73.68      # GroundingPI 11 类等权平均
    astra_score: float = 71.54    # GPT-6 Astra 同口径

    # --- 物理智能 ---
    ood_relative_gain: str = "最高 +24.8%(4 个 OOD 设置全部第一)"
    demo_efficiency: str = "50% 示教数据 SR 28.75% > 基线 75% 数据的 27.75%"
    robotwin_ood: str = "RoboTwin2.0 OOD 17.60%"
    manipulation: str = "6 项操作设置中 5 项成功率第一"

    # 论文按「11 类感知能力」等权平均后仍非第一的单项(沿用正文的如实标注)
    known_caveats: tuple[str, ...] = (
        "VisDrone 40.44 < SenseNova 42.35",
        "TotalText 49.32 < GPT-6 Astra 53.55",
        "Dense200 点定位 81.27 < 86.57",
        "RoboCasa-GR1 Full 37.75% < RynnBrain 39.00%",
    )


CFG = Config()


# ============================================================================
# 1. 模型架构:三段式,但只有「一个」自回归输出头
# ============================================================================
#
# [INSIGHT] 没有检测头 / 分割头 / OCR 头。所有任务都是同一个 next-token
# prediction:位置被量化成 1000 个坐标词元,和文字标签共用词表,
# 因此「定位」在模型眼里只是「继续写一段结构化文本」。

class GroundingPI:
    """视觉编码器 + 投影层 + 语言骨干;输出词表 = 文本词表 ∪ 1000 个坐标词元。"""

    def __init__(self, cfg: Config = CFG):
        self.visual_encoder = build_visual_encoder(cfg.visual_encoder)   # ...
        self.projector = build_projector(cfg.projector)                   # ...
        self.llm = build_llm(cfg.llm_backbone)                            # ...
        self.vocab = extend_vocab_with_coord_tokens(cfg.coord_bins)        # ...

    def forward(self, image: Any, prompt: str) -> Any:
        """图像 -> 视觉特征 -> 投影 -> 与 prompt 拼接 -> 自回归解码 logits。"""
        visual_tokens = self._encode_image(image)        # 1024 projected tokens(论文视觉预算)
        inputs = self._merge(visual_tokens, prompt)
        return self._autoregressive_decode(inputs)       # 一次性输出整段 GAM 串

    def _encode_image(self, image: Any) -> Any: ...
    def _merge(self, visual_tokens: Any, prompt: str) -> Any: ...
    def _autoregressive_decode(self, inputs: Any) -> str: ...


def build_visual_encoder(name: str) -> Any: ...
def build_projector(name: str) -> Any: ...
def build_llm(name: str) -> Any: ...
def extend_vocab_with_coord_tokens(n_bins: int) -> Any: ...


# ============================================================================
# 2. GAM 空间词元协议  —— 本节是真实可运行代码
# ============================================================================
#
# [INSIGHT] 官方 README / 论文 A.1 反复强调:相邻坐标词元之间「不插额外的
# 空格或逗号」(no additional spaces or commas)。服务侧必须设置
# skip_special_tokens=False 与 spaces_between_special_tokens=False,
# 否则解码时词元被拆开、协议作废。上一版文章的代码块就是在这里翻的车。

BOX_START, BOX_END = "<|box_start|>", "<|box_end|>"
REF_START, REF_END = "<|object_ref_start|>", "<|object_ref_end|>"


def quantize(value: float, size: int, bins: int = 1000) -> int:
    """像素坐标 -> 0~999 网格整数。"""
    return max(0, min(bins - 1, round(value / size * (bins - 1))))


def dequantize(token: int, size: int, bins: int = 1000) -> float:
    """0~999 网格整数 -> 像素坐标(线性映射;工程上也可取格子中心 (t+0.5)/bins)。"""
    return token / (bins - 1) * size


def encode_record(label: str, tokens: list[int] | None) -> str:
    """一条预测 -> 一段 GAM 串。tokens 为 None 表示目标不存在,输出 None 拒识。"""
    body = "None" if tokens is None else "".join(f"<{t}>" for t in tokens)
    return f"{REF_START}{label}{REF_END}{BOX_START}{body}{BOX_END}"


def assert_gam_valid(gam: str) -> None:
    """协议守卫:任意两个标记之间不得出现空白或逗号。"""
    if re.search(r">[\s,]+<", gam):
        raise ValueError(
            "GAM 协议被破坏:相邻词元之间插入了空格/逗号。"
            "请确认 skip_special_tokens=False 且 spaces_between_special_tokens=False"
        )


def parse_gam(gam: str) -> list[dict[str, Any]]:
    """GAM 串 -> 结构化预测。坐标个数决定任务粒度:4=框,2=点,None=拒识。"""
    assert_gam_valid(gam)
    pattern = re.escape(REF_START) + r"(.*?)" + re.escape(REF_END) \
              + re.escape(BOX_START) + r"(.*?)" + re.escape(BOX_END)
    out: list[dict[str, Any]] = []
    for label, body in re.findall(pattern, gam, flags=re.S):
        if body.strip() == "None":
            out.append({"label": label, "task": "absent", "coords": None})
            continue
        toks = [int(t) for t in re.findall(r"<(\d+)>", body)]
        if len(toks) == 4:
            task = "bbox"
        elif len(toks) == 2:
            task = "point"
        else:
            task = "invalid"
        out.append({"label": label, "task": task, "coords": toks})
    return out


def to_pixels(record: dict[str, Any], w: int, h: int) -> list[float]:
    """把网格坐标还原成图像像素,供可视化使用。"""
    c = record["coords"]
    if c is None:
        return []
    if record["task"] == "bbox":
        x1, y1, x2, y2 = c
        return [dequantize(x1, w), dequantize(y1, h),
                dequantize(x2, w), dequantize(y2, h)]
    x, y = c
    return [dequantize(x, w), dequantize(y, h)]


# ============================================================================
# 3. 统一任务接口:7 类提示词,塌缩成同一个 next-token 目标
# ============================================================================
#
# [INSIGHT] 换任务只换 prompt,不换 head、不换损失。
# 这也是「范式升级」的落点:一任务一模型 -> 一套生成逻辑通吃所有粒度。

PROMPTS: dict[str, str] = {
    "dense_grounding": "Locate all the instances that match the following categories: car</c>person.",
    "referring_box": "Locate the target referred to by the following description: the red car.",
    "object_point": "Point to: car</c>person.",
    "referring_point": "Point to the target referred to by the following description: the red car.",
    "ocr": "OCR task detect all the text in box format.",
    "layout": "Detect all document layout elements that match the following categories: title</c>text.",
    "gui": "Point to the UI element to click for the following instruction: open the settings menu.",
    "visual_prompt": (
        "Given reference boxes <|box_start|><100><200><500><650><|box_end|> "
        "indicating one or more objects, find all similar objects in the image "
        "and output their bounding boxes."
    ),
}


def build_prompt(task: str, *, phrase: str = "", categories: list[str] | None = None) -> str:
    """按任务拼提示词;视觉提示 / 指代类任务把 phrase 注入模板。"""
    ...


# ============================================================================
# 4. 数据引擎:多教师迭代标注,产出「纯感知」监督
# ============================================================================
#
# [INSIGHT] 由易到难的课程式流水线:先打框,再校验,再修订。
# 关键取舍:训练全程使用无动作标签的纯感知数据,不混带关节动作标签的机器人数据。
# 论文未披露标注图像总量,此处刻意不写具体数字(避免与另一项目串稿)。

SPATIAL_TASK_GROUPS_8 = (
    "detection", "GUI grounding", "referring-expression grounding",
    "referring-expression pointing", "OCR", "document layout",
    "dense pointing and counting", "visual prompting",
)


def data_engine(raw_images: list[Any]) -> list[dict[str, Any]]:
    """Discover -> Localize -> Validate -> Revise / Retire,迭代直到标签稳定。"""
    labels: list[dict[str, Any]] = []
    expert = UnifiedGroundingExpert()          # 被已接受标签持续更新的统一标注专家

    for image in raw_images:
        concepts = discover_concepts(image)                    # 概念发现
        candidates = complementary_teachers(image, concepts)   # 多互补专家定位
        verdict = validate(image, candidates)                  # Pass / Pending / Fail

        if verdict == "Pass":
            labels.append(finalize(candidates))
            expert.update(labels[-1])                          # 接受标签回流训练专家
        elif verdict == "Pending":
            resolved = resolve_by_local_observations(image, candidates, expert)
            labels.append(finalize(resolved))                  # 补监督 + 精化已有标签
        else:
            retire_invalid(candidates)                         # 修订或丢弃无效标签

    apply_field_validation_and_task_dependencies(labels)       # 论文 A.3
    return labels


def discover_concepts(image: Any) -> list[str]: ...
def complementary_teachers(image: Any, concepts: list[str]) -> list[Any]: ...
def validate(image: Any, candidates: list[Any]) -> str: ...
def resolve_by_local_observations(*args: Any) -> list[Any]: ...
def finalize(candidates: list[Any]) -> dict[str, Any]: ...
def retire_invalid(candidates: list[Any]) -> None: ...
def apply_field_validation_and_task_dependencies(labels: list[dict]) -> None: ...
class UnifiedGroundingExpert:
    def update(self, label: dict[str, Any]) -> None: ...


# ============================================================================
# 5. 训练:三大阶段(Pretrain 1 -> Pretrain 2/SFT -> GRPO)
# ============================================================================
#
# [INSIGHT] 量化坐标词元与普通文本词元共享同一 cross-entropy,
# 论文明确「不引入额外的 IoU 回归损失」——这是「一切皆 next-token」最硬的证据。
# SFT: L = -|S|^-1 * sum_{t in S} log p(y_t | I, P, y_<t),S 仅含 assistant 响应。
# GRPO: 每个 prompt 采 G=8 条响应,任务特定奖励(grounding 两个活跃分量 0.7/0.3,
#       OCR 单一复合奖励 1.0),冻结 SFT 参考策略,只更新语言参数。

def train(model: GroundingPI, labels: list[dict[str, Any]]) -> GroundingPI:
    model = pretrain_1_base_vlm(model)          # Stage 1-3
    model = pretrain_2_sft(model, labels)       # Stage 4,坐标对齐
    model = grpo(model)                         # 第三大阶段,强化后训练
    return model


def pretrain_1_base_vlm(model: GroundingPI) -> GroundingPI:
    # Stage 1: 冻结 ViT 与 LLM,只训 projector(image-caption,8192)
    # Stage 2: 全模块,text-only + 图文,LR: LLM 1e-5 / ViT+projector 1e-6
    # Stage 3: 全模块,通用视觉问答 / 指令 / 视频,32768
    # 工程:128 GPUs,BF16,ZeRO-1,FlashAttention,激活重计算,AdamW(0.9,0.95)
    ...


def sft_loss(logits: Any, targets: Any, assistant_mask: Any) -> float:
    """token-level cross-entropy,只在 assistant 响应位置计算(排除 prompt/视觉/pad)。
    语义标签、坐标词元、协议标记一视同仁,无额外 IoU / 回归项。"""
    ...


def pretrain_2_sft(model: GroundingPI, labels: list[dict]) -> GroundingPI:
    for group in SPATIAL_TASK_GROUPS_8:
        batch = sample_group(labels, group)
        loss = sft_loss(model(batch.image, batch.prompt), batch.gam, batch.mask)
        update(model, loss)                     # ...
    return model


def grounding_reward(pred: list[dict], gt: list[dict]) -> float:
    """严格定位 + 输出质量;惩罚重复框、过大框、无效输出。"""
    ...


def ocr_reward(pred: list[dict], gt: list[dict]) -> float:
    """文本-几何一致性复合奖励。"""
    ...


def grpo(model: GroundingPI, group_size: int = CFG.grpo_group) -> GroundingPI:
    """裁剪式 GRPO:eps=0.2,beta=0.02,冻结 SFT 参考,只更新语言参数。"""
    for prompt, gt in rl_prompts():
        samples = [model.generate(prompt) for _ in range(group_size)]   # G=8
        rewards = [composite_reward(s, gt) for s in samples]            # 0.7/0.3 或 1.0
        advantages = standardize_within_prompt(rewards)
        clipped_grpo_step(model, samples, advantages, eps=0.2, beta=0.02)  # ...
    return model


def sample_group(labels: list[dict], group: str) -> Any: ...
def update(model: GroundingPI, loss: float) -> None: ...
def rl_prompts() -> Any: ...
def composite_reward(sample: Any, gt: Any) -> float: ...
def standardize_within_prompt(rewards: list[float]) -> list[float]: ...
def clipped_grpo_step(*args: Any, **kwargs: Any) -> None: ...


# ============================================================================
# 6. 推理:vLLM 服务 + GAM 解析 + 可视化
# ============================================================================

def serve(cfg: Config = CFG) -> str:
    """hf download 权重 -> run.py setup serve -> run.py serve。
    默认 BF16 / TP=1 / 单图 / 16384 上下文 / eager(禁用 CUDA Graph)。"""
    ...


def predict(image: Any, prompt: str, task: str = "bbox", cfg: Config = CFG) -> dict[str, Any]:
    """一次请求 = 一张图 + 一条提示词;返回结构化结果与 valid 标志。"""
    gam = call_vllm(image, prompt, flags=cfg.decode_flags)  # 原始文本,保留特殊词元
    try:
        records = parse_gam(gam)
        valid = all(r["task"] != "invalid" for r in records)
        parse_error = None
    except ValueError as e:
        records, valid, parse_error = [], False, str(e)     # 截断/格式错误保留原始输出
    return {
        "task": task,
        "predictions": [
            {**r, "pixels": to_pixels(r, image.width, image.height)} for r in records
        ],
        "raw_output": gam,
        "valid": valid,
        "parse_error": parse_error,
    }


def call_vllm(image: Any, prompt: str, flags: str) -> str: ...
def visualize(image: Any, result: dict[str, Any]) -> Any: ...


# ============================================================================
# 7. 评测:34 基准 / 44 基线 / 11 类等权平均
# ============================================================================
#
# [INSIGHT] 指标本身也统一:框与 OCR 用 F1_mIoU,点定位用 F1@Point,
# 空间/机器人任务用各自 accuracy。11 类等权平均后 73.68 > 71.54,
# 但这是「综合平均 SOTA」,不是「逐项第一」——caveats 必须随文保留。

METRICS = {
    "bbox / OCR": "F1_mIoU",
    "point": "F1@Point",
    "spatial / robotic": "task-specific accuracy",
}


def evaluate(model: GroundingPI, cfg: Config = CFG) -> dict[str, Any]:
    scores = {name: run_benchmark(model, name) for name in list_34_benchmarks()}
    avg = sum(group_mean(scores, cap) for cap in list_11_capabilities()) / cfg.n_capabilities
    return {
        "groundingpi_avg": round(avg, 2),          # 73.68
        "gpt6_astra_avg": cfg.astra_score,         # 71.54
        "benchmarks": cfg.n_benchmarks,
        "baselines": cfg.n_baselines,
        "caveats": list(cfg.known_caveats),
    }


def list_34_benchmarks() -> list[str]: ...
def list_11_capabilities() -> list[str]: ...
def run_benchmark(model: GroundingPI, name: str) -> float: ...
def group_mean(scores: dict[str, float], capability: str) -> float: ...


# ============================================================================
# 8. 迁移:感知优先 -> 动作学习(物理智能)
# ============================================================================
#
# [INSIGHT] GroundingPI 定位为具身「快系统」的感知原生底座,
# 不替代负责任务规划的通用大模型;下游接 Driving VLA / Robot VLA 的主干对比,
# 论文只证明「更强的感知是更好的物理智能底座」,不是端到端动作策略。

def transfer_to_physical_intelligence(model: GroundingPI, cfg: Config = CFG) -> dict[str, Any]:
    # Driving VLA:图像 + 文本 -> 轨迹词元(x1 y1 ... x6 y6)自回归解码 -> ACT / DiT
    driving = evaluate_driving_vla(model)          # nuScenes 开环 L2 误差
    # Robot VLA:以 GroundingPI 作视觉主干 + 策略头,做 ID / OOD 与数据效率对比
    manipulation = evaluate_robot_vla(model)       # RoboCasa-GR1 / RoboTwin2.0
    return {
        "paradigm": "Basic perception first, action learning next",
        "ood": cfg.ood_relative_gain,
        "data_efficiency": cfg.demo_efficiency,
        "robotwin": cfg.robotwin_ood,
        "manipulation": cfg.manipulation,
        "driving": driving,
        "robot": manipulation,
    }


def evaluate_driving_vla(model: GroundingPI) -> Any: ...
def evaluate_robot_vla(model: GroundingPI) -> Any: ...


# ============================================================================
# 9. 入口:只跑协议自检 + 打印流程骨架,不触发训练
# ============================================================================

PIPELINE = [
    ("1. 架构", "MoonViT-V2 + Projector + Qwen3-4B,单一自回归输出头", "GroundingPI"),
    ("2. 协议", "GAM:1000 个量化坐标词元,相邻词元不插空格/逗号", "parse_gam"),
    ("3. 任务", "7 类提示词塌缩为同一 next-token 目标", "PROMPTS"),
    ("4. 数据", "多教师迭代标注,纯感知数据、无动作标签", "data_engine"),
    ("5. 训练", "Pretrain 1(Stages 1-3) -> SFT(Stage 4) -> GRPO", "train"),
    ("6. 推理", "vLLM 服务 + GAM 解析 + 像素还原", "predict"),
    ("7. 评测", "34 基准 / 44 基线 / 11 类等权平均", "evaluate"),
    ("8. 迁移", "感知优先 -> 动作学习,驱动与机器人 VLA", "transfer_to_physical_intelligence"),
]


def self_check_protocol() -> None:
    """用 README 的三行原始输出格式做一次端到端自检。"""
    print("[协议自检] 原始 GAM 串(严格无空格):")
    gam = "".join([
        encode_record("car", [100, 200, 500, 650]),           # 框 (x1,y1,x2,y2)
        encode_record("car center", [300, 425]),               # 点 (x,y)
        encode_record("absent object", None),                  # 目标缺失 -> None 拒识
    ])
    print("  " + gam)

    records = parse_gam(gam)
    print("\n[协议自检] 解析结果:")
    for r in records:
        print(f"  label={r['label']!r:18s} task={r['task']:8s} coords={r['coords']}")

    print(f"\n[协议自检] 还原到 {CFG.grid_w}x{CFG.grid_h} 像素:")
    for r in records:
        if r["coords"]:
            print(f"  {r['label']:16s} -> {[round(v, 1) for v in to_pixels(r, CFG.grid_w, CFG.grid_h)]}")

    # 反例:故意插入一个空格,确认守卫能拦住
    broken = gam.replace("<100><200>", "<100> <200>")
    try:
        assert_gam_valid(broken)
        print("\n[协议自检] 守卫失效(不应出现)")
    except ValueError as e:
        print(f"\n[协议自检] 反例被正确拦截: {e}")


def main() -> None:
    cfg = CFG
    print("=" * 78)
    print("GroundingPI 全流程伪代码 | 小鹏集团联合北大、港大等高校团队")
    print(f"架构: {cfg.visual_encoder} + {cfg.projector} + {cfg.llm_backbone} ({cfg.params})")
    print("=" * 78)

    print("\n[全流程骨架]")
    for step, desc, symbol in PIPELINE:
        print(f"  {step:10s} {desc:52s} -> {symbol}")

    print()
    self_check_protocol()

    print("\n[关键数字 / 口径核对]")
    print(f"  34 基准 / 44 基线 / 11 类等权平均: {cfg.avg_score} vs GPT-6 Astra {cfg.astra_score}")
    print(f"  Stage 4 训练用的是 {cfg.n_spatial_groups} 个空间任务组(与 11 类评测口径不同)")
    print(f"  物理智能: {cfg.ood_relative_gain}; {cfg.demo_efficiency}")
    print("\n[如实标注:综合平均 SOTA,但以下单项并非第一]")
    for c in cfg.known_caveats:
        print(f"  - {c}")

    print("\n注:建模/训练/评测各节为伪代码(函数体留空),可直接复制进专栏做逐行图解。")
    print("=" * 78)


if __name__ == "__main__":
    main()