Refine documentation and code comments

README.md

@@ -5,7 +5,7 @@ Qwen2.5-0.5B 文本骨干的 VLM 接到缩小版 diffusion expert，验证以下
 训练、保存、重新加载和推理：
 
 1. 图像、自然语言和 ego-relative 历史坐标进入 VLM；
-2. VLM 生成驾驶文字，也可以在未截断验证分支中生成结构化 CoT、meta-action 和
+2. VLM 生成驾驶指令，也可以在未截断验证分支中生成结构化指令、meta-action 和
    离散轨迹 token；
 3. VLM 的 24 层 past-KV 作为 expert 的前缀；
 4. flow-matching diffusion 预测连续 action，再积分回未来 XY 位置。
@@ -60,8 +60,8 @@ dataset，使用前需要在 Hugging Face 接受其许可。
 | `t0_us` | `1503899` |
 | Camera | `camera_front_wide_120fov` |
 | Event cluster | `WORK_ZONES_TEMP_TRAFFIC_CONTROL` |
-| Reasoning source | `reasoning/ood_reasoning.parquet` 中的 `coc` 字段 |
+| 指令来源 | `reasoning/ood_reasoning.parquet` 中的 `coc` 字段 |
-| Reasoning | “Steer left following the temporary lane delineated by traffic cones through the construction zone.” |
+| 驾驶指令 | “Steer left following the temporary lane delineated by traffic cones through the construction zone.” |
 | History | 8 个点，最后一个点为当前 ego-local 原点 |
 | Future source | egomotion 的 64 个未来点，时间步 0.1 s |
 | Expert target | 从 64 点以索引 `[0,4,8,12,16,21,25,29,33,37,42,46,50,54,58,63]` 取 16 点 |
@@ -77,7 +77,7 @@ dataset，使用前需要在 Hugging Face 接受其许可。
 [`scripts/train_stage1_vlm.py`](scripts/train_stage1_vlm.py) 冻结 vision tower、multimodal
 projector 和 LM head，只训练最后 8 个 Qwen2 language layers 及 final norm。当前 checkpoint
 训练 80 steps，共 119,299,968 个可训练参数：loss 从 `3.9338069` 降至
-`0.0000216427`，greedy 输出与数据集 reasoning 逐字一致。
+`0.0000216427`，greedy 输出与数据集驾驶指令逐字一致。
 
 Stage 1 输出到本地 `outputs/stage1_vlm/`，供后续两个分支使用：
 
@@ -99,8 +99,8 @@ checkpoint 继续训练，并添加：
 `[-4, 4]` 的 128 bins，形成 32 个监督 token。
 
 当前结构化 checkpoint 训练 160 steps，loss 从 `14.17349` 降至 `0.0000347825`，最终
-训练 accuracy 为 1.0。未截断 greedy generation 生成 32/32 个正确轨迹 token，其完整
+训练 accuracy 为 1.0。未截断 greedy generation 生成 32/32 个正确轨迹 token。完整
-目标/输出格式为：
+输出格式如下；`cot` 是上游兼容标记，承载的是单句驾驶指令：
 
 ```text
 <|cot_start|>Steer left following the temporary lane delineated by traffic cones through the construction zone.<|cot_end|>
@@ -269,9 +269,9 @@ VLM 的 32 个生成 token 与量化 GT 全部匹配，因此其非零轨迹误
 
 | 字段 | Ground truth | 模型输出 | Exact match |
 | --- | --- | --- | ---: |
-| Reasoning / CoT | Steer left following the temporary lane delineated by traffic cones through the construction zone. | Steer left following the temporary lane delineated by traffic cones through the construction zone. | 是 |
+| 驾驶指令 | Steer left following the temporary lane delineated by traffic cones through the construction zone. | Steer left following the temporary lane delineated by traffic cones through the construction zone. | 是 |
-| Meta-action | `STEER_LEFT`（由 reasoning 规则派生） | `STEER_LEFT` | 是 |
+| Meta-action | `STEER_LEFT`（由驾驶指令规则派生） | `STEER_LEFT` | 是 |
-| Diffusion conditioning text | 同上方 reasoning | 同上方 reasoning | 是 |
+| Diffusion conditioning text | 同上方驾驶指令 | 同上方驾驶指令 | 是 |
 
 ## 本地产物
 
@@ -322,8 +322,8 @@ ruff check src scripts tests
   unicycle acceleration/curvature action space。
 - action mean/scale 由同一个单样本目标轨迹拟合；VLM direct trajectory 的 32 个监督 token
   也来自该目标，因此其近零误差只证明记忆和输出链路。
-- `STEER_LEFT` 是 [`_infer_meta_action`](scripts/train_vlm_trajectory_logits.py) 从 reasoning
+- `STEER_LEFT` 是 [`_infer_meta_action`](scripts/train_vlm_trajectory_logits.py) 从驾驶指令
-  文字规则派生的标签，不是 PAI 中独立提供的 meta-action ground truth。
+  规则派生的标签，不是 PAI 中独立提供的 meta-action ground truth。
 - 当前 expert 用旧 plain-text `"Continuous action:"` prefix 训练，并未与 structured VLM
   prefix 或 exact generation cache 联合训练。structured handoff 证明 24-layer KV 兼容，
   不代表该 prefix 下的轨迹质量已经收敛。

scripts/materialize_pai_subset.py

@@ -46,8 +46,7 @@ def collect_candidates(
                 continue
             if t0_us + int(future_seconds * 1_000_000) > CLIP_DURATION_US:
                 continue
-            # Current PAI releases call the driving rationale/command `coc`.
+            # PAI command text uses `coc`; older metadata shards use `cot`.
-            # Keep `cot` as a compatibility fallback for older metadata shards.
             reasoning = clean_reasoning_text(str(event.get("coc") or event.get("cot") or ""))
             candidates.append({"clip_id": str(clip_id), "t0_us": t0_us, "reasoning": reasoning})
     if not candidates:

src/alpamayo_lite/action_space.py

@@ -13,13 +13,7 @@ class ActionStatistics:
 
 
 class RelativeDeltaActionSpace(nn.Module):
-    """Invertible local validation action space.
+    """Normalized waypoint-delta action space with trajectory conversion."""
-
-    Alpamayo's release model predicts normalized acceleration and curvature. For the
-    architecture smoke test we use normalized per-waypoint XY deltas. It preserves
-    the important contract: diffusion predicts continuous actions, and a deterministic
-    decoder integrates those actions back to relative positions.
-    """
 
     def __init__(self, mean: torch.Tensor | None = None, scale: torch.Tensor | None = None):
         super().__init__()
@@ -39,7 +33,7 @@ class RelativeDeltaActionSpace(nn.Module):
         deltas = future_xy - previous
         flat = deltas.reshape(-1, 2).float()
         mean = flat.mean(dim=0)
-        # A floor keeps a one-clip overfit target numerically well conditioned.
+        # Minimum scale for stable single-clip normalization.
         scale = flat.std(dim=0, unbiased=False).clamp_min(0.05)
         return ActionStatistics(mean=mean, scale=scale)
 

src/alpamayo_lite/expert.py

@@ -151,7 +151,7 @@ class KVCachedDiffusionExpert(nn.Module):
 
         self.architecture = architecture
         self.expert = Qwen2Model(expert_config)
-        # Inputs are always continuous action embeddings, exactly as in Alpamayo.
+        # Continuous action embeddings replace token embeddings.
         self.expert.embed_tokens = None
         self.action_projection = ActionTimeProjection(
             action_dim=action_dim,
@@ -197,9 +197,7 @@ class KVCachedDiffusionExpert(nn.Module):
             device=noisy_action.device,
         )[None, :].expand(batch_size, -1)
         if self.architecture.non_causal_attention:
-            # A prepared 4-D additive mask is returned unchanged by Transformers'
+            # The additive mask exposes the frozen prefix and all diffusion tokens.
-            # mask builder.  This lets all diffusion tokens attend to one another,
-            # while still giving every query access to the frozen VLM prefix KV.
             attention_mask = torch.zeros(
                 batch_size,
                 1,

src/alpamayo_lite/trajectory_tokens.py

@@ -125,7 +125,7 @@ def extract_trajectory_values(
 
 
 class DiscreteTrajectoryMask(LogitsProcessor):
-    """Source-style mask that prevents the VLM from sampling discrete trajectory tokens."""
+    """Mask discrete trajectory-token logits during generation."""
 
     def __init__(self, start_id: int, size: int):
         self.start_id = start_id
@@ -138,7 +138,7 @@ class DiscreteTrajectoryMask(LogitsProcessor):
 
 
 class StopOneTokenAfterBoundary(StoppingCriteria):
-    """Match Alpamayo's handoff behavior so generation updates the cache once after the marker."""
+    """Stop after the token following the handoff boundary."""
 
     def __init__(self, boundary_token_id: int):
         self.boundary_token_id = boundary_token_id

src/alpamayo_lite/vlm.py

@@ -19,10 +19,8 @@ def load_vlm(
     dtype: torch.dtype,
 ) -> tuple[Any, LlavaForConditionalGeneration]:
     processor = AutoProcessor.from_pretrained(model_path, use_fast=False)
-    # The selected third-party checkpoint was saved with this value set to 0,
+    # The checkpoint's CLIP tower produces 576 image tokens after CLS removal.
-    # although its CLIP tower has one CLS token.  LLaVA's "default" feature
+    # num_additional_image_tokens=1 keeps processor expansion aligned with that output.
-    # selection removes CLS after counting it, so the processor must start at 1
-    # to expand <image> into the same 576 tokens produced by the vision tower.
     processor.num_additional_image_tokens = config.num_additional_image_tokens
     model = LlavaForConditionalGeneration.from_pretrained(
         model_path,