VAE-VLA

VAE-VLA 为 Pi0.5 增加有限长度的多模态历史记忆。项目使用 VAE 将从任务开始到当前时刻的 RL token 序列压缩成少量 memory tokens，并将它们送入 Pi0.5 action expert。仓库同时保留了 FM-VLA 的 Force-VAE 复现，因此支持两种微调策略：

当前实现覆盖 RL token 训练、逐帧离线提取、历史 VAE 训练、两种 Pi0.5 微调策略，以及在线推理时的 因果历史队列。详细设计和参数说明见 RL token 历史记忆文档。

方法概览

当前图像 + 语言 + 状态
          |
          v
  冻结的 Pi0.5 VLM 最后一层 embeddings
          |
          v
  RLT Transformer encoder + learnable query
          |
          v
  当前 RL token [2048]
          |
          +----> 从任务开始到当前帧的 RL-token 序列 [T, 2048]
                         |
                         v
               Temporal Latent-Query VAE
                         |
                         v
                 4 个 RL memory tokens --------+
                                                  |
六维 wrench 历史 -> Force-VAE -> 8 个 tokens ----+--> Pi0.5 action expert
short state ---------------------> 1 个 token ----+    （Force 版本）

单帧 RL token 遵循 RLT 的基本结构：在 Pi0.5 VLM 最后一层 prefix embeddings 后追加可学习 query，通过 Transformer encoder 得到紧凑表示；训练时使用自回归 Transformer decoder 重建停止 梯度的原始 embeddings。

历史 VAE 采用 Perceiver-IO 风格的 latent queries。默认将最长 1536 帧的 RL-token 历史压缩为 K=4, dz=256 的 posterior mean；Force-VAE 则将六维 wrench 历史压缩为 K=8, dz=96。两类 latent 经线性投影后追加在 noisy action tokens 后，参与 flow-matching action expert 的预测。

冻结边界

最终微调期间：

• Pi0.5 VLM 冻结；
• RLT encoder、RL-history VAE encoder 和 Force-VAE encoder 冻结；
• Pi0.5 action expert、输入输出投影和新增 memory projector 参与训练。

train_expert_only=true 是强制约束。逐帧 RL token 由冻结 VLM 离线提取；若微调 VLM，缓存表征 会与在线表征发生漂移。策略同样要求 n_action_steps=1，保证每个实际观测都进入历史。

仓库结构

VAE-VLA/
├── data/                                  # LeRobot 原始/派生数据集
├── docs/
│   ├── franka_force.md                    # Force-VAE 数据与训练说明
│   ├── high_dimensional_datasets.md       # 通用高维序列实验说明
│   └── rl_token_history.md                # VAE-VLA 完整设计与训练语义
├── lerobot/src/lerobot/policies/
│   ├── fm_vla/                            # FM-VLA 策略复现
│   └── vae_vla/                           # 两种 VAE-VLA 策略
├── src/vae_vla/
│   ├── cli/                               # 训练、提取与实验入口
│   ├── configs/training/                  # 主链路与复现实验配置
│   ├── force_vae/                         # 通用 Latent-Query VAE
│   └── rl_token_pipeline.py               # 冻结 Pi0.5 embedding 提取
├── outputs/                               # 本地 checkpoint 与训练输出
└── papers/                                # 参考论文

主要实现：

• VAE-VLA 模型与在线队列
• 策略配置
• RLT encoder/decoder
• 冻结的历史 VAE
• 通用 Latent-Query VAE

环境准备

项目要求 Python 3.12 或更高版本，使用 uv 管理根项目和内嵌 LeRobot workspace：

uv sync --frozen

pyproject.toml 默认使用 CUDA 12.8 的 PyTorch index。若服务器驱动不兼容，应先调整 PyTorch index 并重新生成 lockfile，不要复制其他机器的 .venv。

完整训练前需要以下输入：

data/lerobot_dataset_v2/                  # 原始 LeRobot 数据集
outputs/pi05/pretrained_model/            # 用于 RLT 和下游策略的同一 Pi0.5 checkpoint
outputs/VAE/t1536_k8/best.pt              # 仅联合 Force 版本需要

RL token 提取会创建 data/lerobot_dataset_v2_rl_tokens，不会修改原数据集。LeRobot 特征修改工具 会复制视频，因此需要为派生数据集预留接近原数据集大小的额外空间。

端到端训练

以下步骤存在严格依赖关系，应按顺序执行。

1. 训练当前观测 RL token

uv run vae-vla-train-rl-token \
  --dataset-root data/lerobot_dataset_v2 \
  --repo-id local/vae_vla_source \
  --pi05-checkpoint outputs/pi05/pretrained_model \
  --output outputs/rl_token/pretrained_model \
  --device cuda

默认对完整视觉语言 prefix 做 mask-aware 64-token pooling，再训练 RLT 重建目标。输出 checkpoint：

outputs/rl_token/pretrained_model/best.pt

2. 提取整个数据集的逐帧 RL token

uv run vae-vla-extract-rl-tokens \
  --dataset-root data/lerobot_dataset_v2 \
  --repo-id local/vae_vla_source \
  --pi05-checkpoint outputs/pi05/pretrained_model \
  --rl-token-checkpoint outputs/rl_token/pretrained_model/best.pt \
  --output-root data/lerobot_dataset_v2_rl_tokens \
  --output-repo-id local/vae_vla_rl_tokens \
  --device cuda

派生数据集新增 observation.rl_token: float32[2048]，并保留原始图像、状态、动作、任务和 episode 元数据。

3. 训练 RL-token 历史 VAE

uv run vae-vla-train-history-vae \
  --config src/vae_vla/configs/training/rl_history_k4.json \
  --output outputs/rl_history_vae \
  --device cuda

数据按完整 episode 划分。训练样本随机选择“任务起点到当前帧”的因果前缀，均值和标准差只由 训练 episode 计算并保存进 checkpoint：

outputs/rl_history_vae/best.pt

4. 微调 Pi0.5

仅使用 RL 历史：

uv run lerobot-train \
  --config_path=src/vae_vla/configs/training/vae_vla_franka.yaml

同时使用 RL 历史和 Force 历史：

uv run lerobot-train \
  --config_path=src/vae_vla/configs/training/vae_vla_force_franka.yaml

两个版本共享同一个 RL-token 派生数据集和历史 VAE checkpoint。联合版额外读取现有 outputs/VAE/t1536_k8/best.pt，并继承 FM-VLA 的 force EMA、独立归一化和 short-state 路径。

关键配置

后两份复现配置保留了早期实验的数据和输出目录约定；重新运行前应根据本机数据位置检查其中的 dataset.path、dataset.root 和 force_vae_checkpoint。VAE-VLA 主链路使用前三份配置。

修改实验参数时必须保持以下契约：

1. RLT 训练、离线提取和最终策略必须使用同一个 Pi0.5 VLM checkpoint。
2. pool_size、image_only、RL token 维度和 Transformer 结构必须与 RLT checkpoint 一致。
3. rl_history_length 必须与历史 VAE 的 sequence_length 容量一致。
4. episode 长度超过容量时不会静默裁掉任务起点；应增大长度并重新训练历史 VAE。
5. observation.rl_token 保持原始 embedding 空间，只应用历史 VAE checkpoint 自带的 mean/std，不参与 Pi0.5 的 state quantile normalization。

在线推理语义

调用 policy.reset() 会清空 RL-token 和 Force 历史。每个控制步执行一次当前观测的 VLM prefix 前向，生成当前 RL token，将其追加到队列后再压缩历史。训练和部署因此都使用“任务开始到当前帧” 的因果上下文。

历史超过 1536 帧时，队列丢弃最早 token，退化为固定长度滑动窗口。当前实现不是无限流式记忆； 更长任务需要增大容量，或进一步实现分块、层级 latent、递归状态更新。

测试

VAE-VLA、数据历史采样和 FM-VLA 回归测试：

PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 uv run --with pytest pytest -q \
  lerobot/tests/policies/vae_vla/test_vae_vla.py \
  tests/test_data.py \
  lerobot/tests/policies/fm_vla/test_fm_vla.py

通用 VAE 的轻量训练 smoke test：

uv run vae-vla-train-force \
  --config src/vae_vla/configs/training/smoke.json \
  --output outputs/smoke \
  --epochs 2 \
  --max-steps 2

附加 VAE 实验

src/vae_vla/force_vae 也可独立用于长时间序列压缩研究，目前支持：

• 合成稀疏事件序列；
• ETTh1 电力变压器序列；
• PEMS-SF 高维交通传感器序列；
• Moving MNIST 视频序列；
• Franka 六维 wrench episode。

这些实验用于比较不同序列长度、输入维度和 latent 数量下的率失真、频域误差、下游 probe 与 编码延迟，不是运行 VAE-VLA 主链路的前置条件。详情见 高维数据集说明和 Franka Force-VAE 指南。

当前边界

• 完整 RLT 训练、全数据集 embedding 提取和 Pi0.5 微调都是 GPU 离线任务，不会在安装时自动执行。
• RLT 派生数据依赖固定 VLM；更换或微调 VLM 后必须重新提取 RL token 并重训历史 VAE。
• VAE 重建误差只能衡量信息保真度，最终是否获得有效任务记忆仍需用下游成功率、长程依赖任务和 RL/模仿学习对照实验验证。
• 当前配置面向单卡训练；多卡机器更适合并行运行不同 K、历史长度或策略版本的独立实验。

参考

• RLT: Bootstrapping RL with Vision-Language-Action Models
• FM-VLA 项目页
• Physical Intelligence Pi0.5
• Perceiver / Perceiver IO

傅里叶变换/小波变换

koopman时间序列压缩与预测

FTDASC