YuanhaoXD/Lam_VGGT

LaWAM Stage 1 · VGGT-1B 版本

LaWAM Stage 1（latent action model + latent world model）的复现，把原本冻结的 DINOv3 视觉编码器换成了冻结的 VGGT-1B（3D 几何基础模型，2048 维）。 在 LIBERO 上训练 40 epoch。前向 decoder 预测的是未来观测的特征，不是像素。

视觉编码器	facebook/VGGT-1B，冻结，patch-14 @ 518 → 37×37 池化到 16×16 = 256 token
训练数据	jialei02/libero_merged_no_noops_20hz（LeRobot v3.0）
检查点	epoch 39 / global_step 42720，fp32
Stage 2 只需	weights/lawm_decoder.safetensors（923.7 MB / 230.9 M 参数）
关键结果	val recon 0.130 → 0.0561，cos_sim 0.992，z2action_r2 = 0.5523

完整训练配置、评估、以及一个反转了早期乐观结论的动作热力图分析见 LAWAM_VGGT_STAGE1_REPORT.md §11.3 —— 简言之：latent action z 控制的是一块固定的工作区区域，并不跟着机械臂走。

许可

CC-BY-NC-4.0（非商用）。上游 RLinf/LaWAM 的代码沿用 StarVLA 的 MIT，但本仓库的权重是在冻结的 VGGT-1B （cc-by-nc-4.0）特征上训练得到的，因此按更严格的那一方标注。 VGGT-1B 权重本身不在本仓库内，使用者自行从 facebook/VGGT-1B 拉取。

---

目录结构

├── code/                              ★本次全部代码改动（264 KB）
│   ├── changes_vs_upstream.patch      333 行，对 3 个已有文件的 diff
│   ├── latent_action_model/core/      vjepa_encoder / lam_lightinng / lam_model
│   ├── latent_action_model/config/    vggt_vae.yaml（新增）
│   └── tools/                         10 个新增脚本
├── weights/                           2.58 GB
│   ├── lawm_decoder.safetensors       923.7 MB  230.9M 参数  ← Stage 2 只要这个
│   ├── lam_idm_encoder.safetensors   1243.0 MB  310.7M 参数
│   ├── state_head.safetensors         411.6 MB  102.9M 参数
│   ├── misc.safetensors                 9.2 KB  VAE 的 mu/logvar/pre_norm
│   └── manifest.json                  含每个文件的 sha256
├── LAWAM_VGGT_STAGE1_REPORT.md        主报告（28 KB，含 §11 可视化解读）
├── LAWAM_VGGT_PROBE_ANALYSIS.md       探针分析（20 KB）
├── docs/figs/                         主报告内嵌的 5 张图（6 MB，勿改名）
├── figures/                           训练曲线 + curves.csv（323 KB）
├── logs/                              train.log.gz + TB event（1.6 MB）
├── configs/                           三份 yaml
└── launch_vggt_lam.sh

code/ 是本次工作的全部代码产出：3 个改动文件（另附 patch）+ 1 个新 config + 10 个新脚本。 上游 fork 的其余部分未改动，不在发布件内。

---

模型多大

	大小	说明
训练 ckpt epoch=39.ckpt	7.74 GB	其中 5.16 GB 是 AdamW 优化器状态
→ 剥掉优化器后的权重	2.58 GB	449 个张量，fp32，644.6 M 参数
→ Stage 2 实际需要	923.7 MB	只有 decoder；IDM 训完即弃

原始 ckpt 已经是剥掉冻结 VGGT 之后的（on_save_checkpoint 钩子）， 否则会是 11 GB。VGGT-1B（909 M，冻结）不在本仓库里 —— 它是未经修改的第三方权重， 使用者从 facebook/VGGT-1B 自行拉取即可，重传一遍没有意义。

导出与验证

python tools/export_release_weights.py     # 7.74 GB → 2.58 GB
python tools/verify_release_weights.py     # 上传前必跑

验证结果：449/449 张量与原 ckpt 逐位一致，0 缺失，0 数值偏差。 本仓库 4 个 safetensors 的 sha256 也已与 manifest.json 逐一核对通过。

⚠️ safetensors 用 mmap 写盘，在 root-squash NFS 上会直接 EPERM。 导出脚本因此先写 /dev/shm 再 copyfile。

---

训练收敛曲线

figures/training_curves.png，四个面板，数据同时来自 TensorBoard event（val，逐 epoch）和 40 MB 训练日志（train，42721 个点， 从 tqdm 进度条正则抽出——Lightning 没把 train 指标写进 TB）。 figures/curves.csv 是逐 epoch 的数值表。

面板	内容
左上	recon loss 对数轴，train + val 叠加，标注最低点 ep24 = 0.0526
右上	val 线性轴，高亮 ep24 后的震荡带
左下	train/val 并列 + 比值曲线（过拟合检查）
右下	cos_sim 0.970→0.992、state_loss 降 56×

关键数字：val 0.130 → 0.0561，train/val gap 1.43×，全程无发散。

⚠️ 训练期 val 集只有 109 个样本（val_tail_ratio: 0.001），±0.005 纯属噪声。 真正可信的评估用的是 5% 留出（84 episodes / 640 样本），见主报告 §4。

---

关于上游仓库

RLinf/LaWAM 的 clone 有 15 GB（大部分是数据和 ckpt），本次真正的改动只有：

 latent_action_model/core/vjepa_encoder.py   | 248 +++   ← VGGTEncoder 主体
 latent_action_model/core/lam_lightinng.py   |  28 ++    ← ckpt 瘦身钩子
 latent_action_model/core/lam_model.py       |  10 ++
 latent_action_model/config/vggt_vae.yaml    | 新增
 tools/*.py                                  | 新增 10 个
 LAWAM_VGGT_STAGE1_REPORT.md                 | 新增
 LAWAM_VGGT_PROBE_ANALYSIS.md                | 新增
 docs/figs/*.png                             | 新增 5 张

想复现代码改动：fork RLinf/LaWAM → 开分支 → 把本目录 code/ 下的文件按原路径放回。 code/changes_vs_upstream.patch 可直接 git apply 到干净的上游 clone 上复现那 3 个文件的改动。

路径映射：

code/latent_action_model/core/*.py    → latent_action_model/core/
code/latent_action_model/config/*.yaml→ latent_action_model/config/
code/tools/*.py                       → tools/
LAWAM_VGGT_*.md                       → 仓库根
docs/figs/*.png                       → docs/figs/   （报告内嵌图，勿改名）

注：上游 latent_action_model/train.sh 里有一个明文 W&B key（上游自己的，非本次引入； 本次训练全程 WANDB_MODE=disabled）。基于上游二次分发时记得替换成环境变量占位符。 本发布件已扫过，不含该 key。

---

怎么用

hf download YuanhaoXD/Lam_VGGT weights/lawm_decoder.safetensors --local-dir .

from safetensors.torch import load_file

# Stage 2 只需要 decoder。键名保留 "decoder." 前缀（只剥掉了 Lightning 的 "lam." 一层），
# 所以是按整个 LaM 模块的命名空间来的，装载时注意对齐。
sd = load_file("weights/lawm_decoder.safetensors")
print(len(sd), "tensors,", sum(v.numel() for v in sd.values()) / 1e6, "M params")
# -> 230.9 M，键形如 decoder.blocks.0.attn.qkv.weight

# 只要 decoder 子模块本身：
sub = {k[len("decoder."):]: v for k, v in sd.items()}

# 冻结的视觉塔要自己拉，本仓库不含（未改动的第三方权重，重传无意义）
# from vggt.models.vggt import VGGT; VGGT.from_pretrained("facebook/VGGT-1B")

weights/manifest.json 记录了每个文件的 sha256 和对应的 prefix，可校验完整性。 输入侧的约定（S 必须为 1、16×16=256 token、特征必须标准化）见报告 §11 和 code/latent_action_model/core/vjepa_encoder.py 里的 VGGTEncoder。

---

复现

# 训练（8×A100-80G，约 56 h）
./launch_vggt_lam.sh

# 评估三探针
python tools/eval_lam_probes.py --n-batches 40

# 出图
python tools/viz_vggt_lam.py --n-samples 6 --seed 0
python tools/viz_action_heatmap.py --batch 16 --n-show 6 --shuffle
python tools/plot_training_curves.py

数据：jialei02/libero_merged_no_noops_20hz（HF，LeRobot v3.0，1.97 GB）。 本地那份 LIBERO 是 v2.1，官方 fork 不认。国内走 HF_ENDPOINT=https://hf-mirror.com。

---

引用

本工作基于 RLinf/LaWAM（LaWM / latent action model） 与 VGGT（arXiv:2503.11651）。 使用本仓库权重时请同时引用这两项原始工作。