3DThinker

Overview概述

3DThinker is the first framework that enables VLMs to perform "3D mentaling" during reasoning — imagining 3D spatial structures from limited-view images without any 3D prior input. Training consists of two stages: Stage 1 aligns the VLM's 3D latent with a 3D foundation model (VGGT) via supervised training; Stage 2 refines the entire reasoning trajectory through GRPO reinforcement learning using outcome-based rewards.

3DThinker 是首个让 VLM 在推理过程中进行"3D 心象"的框架——无需任何 3D 先验输入，即可从有限视角图像中想象 3D 空间结构。训练分两阶段：阶段 1 通过监督训练将 VLM 的 3D 隐表示与 3D 基础模型（VGGT）对齐；阶段 2 通过 GRPO 强化学习，仅基于结果信号优化完整推理轨迹。

Data Preprocessing数据预处理

VGGT Feature ExtractionVGGT 特征提取

Download VGGT-1B weights into models/. The script processes each scene's images through VGGT and saves aggregated 3D tokens as .npz files with shape [N, 1374, 2048].

下载 VGGT-1B 权重到 models/。脚本将每个场景的图像通过 VGGT 处理，保存聚合的 3D token 为 .npz 文件，shape 为 [N, 1374, 2048]。

from vggt.models.vggt import VGGT

# Load VGGT-1B model
model = VGGT.from_pretrained("models/VGGT-1B")
model.eval().cuda()

# For each training sample: load N views, extract features
images = load_and_preprocess(image_paths)      # [N, 3, H, W]
with torch.no_grad():
    # VGGT outputs aggregated 3D tokens per view
    features = model.forward_features(images)  # [N, P_3D=1374, C=2048]

# Save features for downstream training
np.savez(output_path, feature=features.cpu().numpy())
# Output: data/feature_vggt/{idx}/vggt.npz
# Also saves resized images: data/resized_images/

CoT Data GenerationCoT 数据生成

Chain-of-thought reasoning traces are generated using GPT-4.1. Each trace starts with <output_3D> — the placeholder for the model's mental 3D scene — followed by structured thinking and answer blocks.

使用 GPT-4.1 生成思维链推理轨迹。每条轨迹以 <output_3D> 开头——作为模型"3D 心象"的占位符——随后是结构化的思考和回答块。

SYSTEM_PROMPT = """You are a helpful assistant that can answer questions
about 3D spatial relationships from images.

The mental 3D scene you imagine will be represented by the special token
<output_3D>, which should be presented at the beginning of your response.

Your response should follow this format:
<output_3D><think> [step-by-step spatial reasoning] </think>
<answer> [final answer letter] </answer>"""

# Generate CoT with GPT-4.1 for each QA pair
response = client.chat.completions.create(
    model="gpt-4.1",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": [
            {"type": "image_url", "image_url": {"url": img_url}},
            {"type": "text",      "text": question}
        ]}
    ]
)

Data Cleaning Pipeline数据清洗流程

Four scripts run sequentially to produce the final training JSONL:

四个脚本顺序执行，生成最终训练用的 JSONL 文件：

# 1. Generate chain-of-thought reasoning data
python preprocessing/produce_cot.py

# 2. Remove non-compliant entries (e.g., missing </output> tags)
python preprocessing/clean.py

# 3. Filter low-quality / useless data
python preprocessing/remove.py

# 4. Match VGGT feature indices to training samples
python preprocessing/jsonl_add_idx.py

# Output: data/data_output3d_begin_10k_resized.jsonl

Each entry in the final JSONL contains:

最终 JSONL 的每条数据包含：

{
  "text_input": "Question about spatial relationship...",
  "image_input": ["img1.jpg", "img2.jpg", "img3.jpg", "img4.jpg"],
  "text_output": "<output_3D><think>Looking at the images, I can see...</think>\n<answer>B</answer>",
  "idx": 42
}

Stage 2: Reinforced Trajectory Refinement (GRPO)阶段 2：强化推理轨迹优化 (GRPO)

Stage 2 optimizes the full reasoning trajectory using Group Relative Policy Optimization (GRPO). The model generates multiple candidate responses, scores them with three reward signals, and updates the policy. No labeled 3D data is used — only outcome signals.

阶段 2 使用 Group Relative Policy Optimization (GRPO) 优化完整推理轨迹。模型生成多个候选回复，通过三个奖励信号评分，更新策略。不使用标注的 3D 数据——仅基于结果信号。

Reward Functions奖励函数

Three reward functions score each generated trajectory. Each reward returns a scalar in [0, 1] per sample; GRPO then computes group-relative advantages over the num_generations=8 rollouts and updates the LoRA weights with a KL penalty (β=0.04) against the reference policy.

三个奖励函数为每条生成的轨迹评分。每个 reward 输出 [0, 1] 区间的标量；GRPO 随后在 num_generations=8 条采样上计算组内相对优势（以均值为基线），并带 KL 惩罚（β=0.04）更新 LoRA 权重。

1. Format reward — a regex check on the full trajectory. The completion must contain the latent block followed by <think>…</think> and then <answer>…</answer>. Returns 1.0 if the structure matches, 0.0 otherwise. This keeps the model from collapsing to plain text and forces it to actually emit the 12 latent tokens that the projector will read.
   Example — same MindCube question, two rollouts:
   rollout #3  →  format = 1.0

   <|latent_start|>⟨z₁⟩⟨z₂⟩…⟨z₁₂⟩<|latent_end|>
   <think>The camera rotates 90° clockwise, so the
   chair that was on the left should now be in front…</think>
   <answer>B</answer>

   rollout #5  →  format = 0.0

   The chair that was on the left moves to the front
   after a 90° rotation, so the answer is B.

   // no latent block, no <think>/<answer> tags → regex misses → 0.0 even though the answer is correct.
2. Response (answer) reward — the content inside <answer> is parsed for the option letter (A–E) via priority-based regex, then compared with the ground-truth letter. Exact match → 1.0, otherwise 0.0. This is the only outcome signal that ties the trajectory to the actual benchmark task.
   Example — ground-truth letter is B:
   extract_answer( … )extract_option(…)==r

   "B"B✓1.0

   "The answer is (B)."B✓1.0

   "D"D✗0.0

   "probably the chair"None✗0.0
   Note: the priority regex only fires on a clean letter — verbose answers without an extractable option still score 0.0.
3. Latent-similarity reward — the 12 latent token hidden states from the rollout are fed through the (frozen) 3D Projector together with the question's image patch embeddings, producing a predicted [N, 1374, 2048] feature map. We compare it against the pre-computed VGGT ground-truth features for the same scene via mean cosine similarity, and remap from [-1, 1] to [0, 1]. This rewards trajectories whose internal 3D "mental image" stays geometrically faithful, even when the answer letter alone wouldn't reveal that.
   Example — same question (idx mc_00471), two rollouts that both answer B correctly:
   # load pre-computed teacher features for this scene
   vggt = np.load("../../data/feature_vggt/mc_00471/vggt.npz")['feature']  # shape [1, 1374, 2048]
   
   rollout #3:  projector( z₁…z₁₂, img_emb ) → pred  ⟶  mean_cos( pred, vggt ) =  +0.71  ⟶  (0.71 + 1)/2 =  0.855
   rollout #7:  projector( z₁…z₁₂, img_emb ) → pred  ⟶  mean_cos( pred, vggt ) =  −0.04  ⟶  (−0.04 + 1)/2 =  0.480
   Both rollouts win the response reward (answer = B), but only rollout #3 also wins here — its latent tokens encode the rotated layout faithfully, while rollout #7 happened to land on B by chance with a noisy 3D representation.

1. Format reward（格式奖励） — 对整条轨迹做一次正则检查：输出必须包含 latent 段，随后是 <think>…</think>，再接 <answer>…</answer>。结构命中返回 1.0，否则 0.0。作用是防止模型倒退为纯文本、并强迫它真的吐出那 12 个供投影器读取的 latent token。
   例子 —— 同一道 MindCube 题，两次采样：
   采样 #3  →  format = 1.0

   <|latent_start|>⟨z₁⟩⟨z₂⟩…⟨z₁₂⟩<|latent_end|>
   <think>相机顺时针旋转 90°，原本在左边的椅子
   现在应该出现在正前方……</think>
   <answer>B</answer>

   采样 #5  →  format = 0.0

   原本在左边的椅子在 90° 旋转之后会到正前方，
   所以答案是 B。

   // 没有 latent 段，也没有 <think>/<answer> 标签 → 正则匹配失败 → 即使答案对了也只能拿 0.0。
2. Response reward（答案奖励） — 对 <answer> 内部的内容用优先级正则抽出选项字母（A–E），与真值字母严格比对：完全一致 → 1.0，否则 0.0。这是唯一直接把轨迹绑定到下游任务的 outcome 信号。
   例子 —— 真值是 B：
   extract_answer( … )extract_option(…)==r

   "B"B✓1.0

   "答案是 (B)。"B✓1.0

   "D"D✗0.0

   "大概是椅子吧"None✗0.0
   注意：优先级正则只认得「干净的字母」——含糊其辞、抽不出选项的回答即使语义上接近也只能记 0.0。
3. Latent-similarity reward（latent 相似度奖励） — 把采样轨迹中 12 个 latent token 的隐状态与该问题的图像 patch 嵌入一起送进冻结的《3D Projector》，得到预测的 [N, 1374, 2048] 特征，再与预计算好的 VGGT ground-truth 特征取平均余弦相似度，最后将 [-1, 1] 映射到 [0, 1]。它奖励那些内部「心象 3D」仍然几何上忠实的轨迹 — 这是单看选项字母看不出来的。
   例子 —— 同一道题（idx mc_00471），两条都答对 B 的采样：
   # 加载这道题预先算好的 VGGT 教师特征
   vggt = np.load("../../data/feature_vggt/mc_00471/vggt.npz")['feature']  # 形状 [1, 1374, 2048]
   
   采样 #3：  projector( z₁…z₁₂, img_emb ) → pred  ⟶  mean_cos( pred, vggt ) =  +0.71  ⟶  (0.71 + 1)/2 =  0.855
   采样 #7：  projector( z₁…z₁₂, img_emb ) → pred  ⟶  mean_cos( pred, vggt ) =  −0.04  ⟶  (−0.04 + 1)/2 =  0.480
   两条轨迹都拿到了 response reward（答案都是 B），但只有 #3 在这里也拿了高分——它的 latent token 真的把旋转后的布局编码对了；#7 只是「蒙对」字母，内部 3D 表示其实是噪声。

In the launch script, all three are switched on simultaneously via --reward_funcs format response latent_sim, and they are summed with equal weight per rollout (no per-reward coefficient). The format reward acts as a hard gate — a malformed trajectory loses 1 point regardless of accuracy — while the answer and latent rewards push the policy toward both correct and geometrically grounded reasoning.

启动脚本中三项同时打开（--reward_funcs format response latent_sim），等权相加（没有单独的 reward 系数）。Format 奖励充当硬门槛——格式出错不论答案对不对都会扣 1 分；Response 和 Latent 则共同推动策略同时走向「答案正确」与「几何上不脱靠」两个方向。

@staticmethod
def format_reward(completions, **kwargs):
    """Reward 1: Format compliance — full trajectory must match pattern."""
    # Regex: latent tokens → <think>...</think> → <answer>...</answer>
    pattern = r".*?<\|latent_start\|>.*?<think>.*?</think>\s*<answer>.*?</answer>"

    completion_contents = [c[0]["content"] for c in completions]
    matches = [
        re.search(pattern, content, re.DOTALL) is not None
        for content in completion_contents
    ]
    rewards = [1.0 if match else 0.0 for match in matches]
    return rewards

def extract_answer(xml_text):
    """Extract content between <answer> and </answer> tags."""
    match = re.search(r'<answer>\n?(.*?)\n?</answer>', xml_text, re.DOTALL)
    return match.group(1).strip() if match else ""

@staticmethod
def response_text_reward(completions, answer, prompts, **kwargs):
    """Reward 2: Answer accuracy — option letter matching."""
    contents = [c[0]["content"] for c in completions]
    rewards = []
    for content, gt_response in zip(contents, answer):
        pred_response = extract_answer(content)
        # extract_option: priority-based regex to find A-E letter
        if extract_option(pred_response) and \
           extract_option(gt_response) == extract_option(pred_response):
            rewards.append(1.0)
        else:
            rewards.append(0.0)
    return rewards

@staticmethod
def sim_reward(index, prompts, image_embeddings, extracted_emb, **kwargs):
    """Reward 3: 3D latent similarity — cosine sim with VGGT features."""
    rewards = []
    for idx, prompt in zip(index, prompts):
        target_dtype = torch.float32
        # Run the projector on generated latent embeddings
        decoder_feature = projector_model(
            extracted_emb.to(target_dtype).to("cuda:0"),
            image_embeddings.to(target_dtype).to("cuda:0")
        )
        decoder_feature_norm = decoder_feature / decoder_feature.norm(
            dim=-1, p=2, keepdim=True
        )

        # Load pre-computed VGGT ground truth features
        data = np.load(f'../../data/feature_vggt/{idx}/vggt.npz')
        feature_3d = torch.tensor(data['feature']).to(
            device=decoder_feature.device, dtype=decoder_feature.dtype
        ).squeeze()
        feature_3d_norm = feature_3d / feature_3d.norm(
            dim=-1, p=2, keepdim=True
        )

        # Cosine similarity → [0, 1] reward
        cos_sim = F.cosine_similarity(
            decoder_feature_norm, feature_3d_norm
        ).mean()
        reward = (cos_sim + 1) / 2  # map [-1,1] → [0,1]
        rewards.append(reward)
    return rewards

Training Launch & Weight Merging训练启动与权重合并

# Stage 2: GRPO Reinforcement Learning
conda activate 3DThinker-stage2

# Path to Stage 1 checkpoint
model_path="../../models/3DThinker-S1-Qwen2.5-VL-3B_mlp6_lr1e-4_latent12"
EXP_NAME="3DThinker-S1+S2-Qwen2.5-VL-3B_lr1e-5_gen8"

torchrun --nproc_per_node=8 \
    --nnodes=1 --node_rank=0 \
    --master_addr=127.0.0.1 --master_port=12346 \
  src/open_r1/grpo_jsonl.py \
    --use_vllm True \
    --model_name_or_path $model_path \
    --output_dir checkpoints/rl/${EXP_NAME} \
    --reward_funcs format response latent_sim \
    --beta 0.04 \
    --num_generations 8 \
    --max_completion_length 2048 \
    --use_peft true \
    --lora_r 64 \
    --lora_alpha 128 \
    --lora_dropout 0.05 \
    --lora_task_type CAUSAL_LM \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --learning_rate 1e-5 \
    --num_train_epochs 1 \
    --freeze_vision_modules true \
    --deepspeed local_scripts/zero3.json \
    --bf16

# Merge LoRA adapters into the base model
conda activate SFT
llamafactory-cli export merge.yaml

# merge.yaml specifies:
#   model_name_or_path: checkpoints/stage1    (base)
#   adapter_name_or_path: checkpoints/stage2  (LoRA)
#   export_dir: checkpoints/merged            (output)

Citation引用

@inproceedings{chen2026think,
  title   = {Think with 3D: Geometric Imagination Grounded
             Spatial Reasoning from Limited Views},
  author  = {Chen, Zhangquan and Zhang, Manyuan and Yu, Xinlei
             and Luo, Xufang and Sun, Mingze and Pan, Zihao
             and An, Xiang and Feng, Yan and Pei, Peng
             and Cai, Xunliang and Huang, Ruqi},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year    = {2026}
}