VLA Framework: Vision-Language-Action Policy Framework

VLA Framewrok 是 roboimi 生态系统中的下一代具身智能策略框架。它采用完全解耦与基于组合的架构设计，支持视觉语言模型（VLM）、投影层（Projector）与动作生成头（Action Head）的灵活搭配。

本框架基于 Hydra 进行配置管理，并采用 HDF5 作为标准数据格式。

---

🏗 架构概览 (Directory Structure)

我们采用“接口与实现分离”以及“代码与配置镜像映射”的设计原则。

roboimi/vla/
├── agent.py                # [Core] VLAAgent 组装类，负责串联各个模块
├── conf/                   # [Config] Hydra 配置文件 (单一真值源)
│   ├── config.yaml         # 主入口配置
│   ├── agent/              # Agent 结构定义 (定义模块间的连接与插值)
│   ├── backbone/           # 视觉骨干配置 (e.g., SigLIP, CLIP)
│   ├── projector/          # 投影层配置 (e.g., MLP, Perceiver)
│   ├── head/               # 动作头配置 (e.g., Diffusion, ACT)
│   └── data/               # 数据流配置
├── core/                   # [Interface] 抽象基类
│   ├── base_vlm.py         # VLMBackbone (ABC)
│   └── base_policy.py      # ActionHead (ABC)
├── models/                 # [Implementation] 具体模型实现
│   ├── backbones/          # 视觉模型 (Sub-package)
│   ├── projectors/         # 投影层 (Sub-package)
│   └── heads/              # 策略头 (Sub-package)
├── data/                   # [Data Pipeline] Dataset 与 DataLoader
├── modules/                # [Building Blocks] 通用组件 (Encoders, Fusion)
└── scripts/                # [Utilities] 数据转换与维护脚本

---

🚀 快速开始 (Quick Start)

1. 环境依赖

请确保安装以下核心库：

pip install hydra-core h5py zarr diffusers transformers

2. 启动训练 (Training)

训练入口脚本通常位于 demos/vla_scripts/train_vla.py。 由于使用了 Hydra，您可以在命令行动态组合模型架构：

# 1. 默认训练 (SigLIP + MLP + Diffusion)
python demos/vla_scripts/train_vla.py

# 2. 切换视觉骨干为 CLIP
python demos/vla_scripts/train_vla.py agent/backbone=clip

# 3. 切换投影层为 Perceiver Resampler
python demos/vla_scripts/train_vla.py agent/projector=perceiver

# 4. 修改超参数 (例如 batch size)
python demos/vla_scripts/train_vla.py train.batch_size=32

# 5. 调试模式 (使用 Tiny 模型快速跑通流程)
python demos/vla_scripts/train_vla.py agent=tiny

---

🛠 开发指南 (Developer Guide)

1. 添加新的视觉骨干 (New Backbone)

1. 代码: 在 models/backbones/ 下新建文件 (如 my_model.py)，继承 VLMBackbone。
2. 导出: 在 models/backbones/__init__.py 中添加导出。
3. 配置: 在 conf/backbone/ 下新建 my_model.yaml。
   • 注意: 必须定义 output_dim，供 Projector 引用。

2. 添加新的投影层 (New Projector)

Projector 负责将 VLM 特征维度对齐到 Agent 的 Embedding 维度。

1. 代码: 在 models/projectors/ 下实现 nn.Module。
2. 配置: 在 conf/projector/ 下新建 YAML 文件。
   • 关键: 设置 input_dim: ??? 和 output_dim: ???，让 Hydra 在 agent/default.yaml 中自动插值填充。

3. 添加新的动作头 (New Action Head)

1. 代码: 在 models/heads/ 下新建文件，继承 ActionHead。
   • 必须实现 compute_loss(context, actions) 和 predict_action(context)。
2. 配置: 在 conf/head/ 下新建 YAML 文件。
   • 同样建议设置 input_dim: ??? 以保持动态性。

---

📊 数据流水线 (Data Pipeline)

本框架强制使用 HDF5 格式以优化 IO 性能。

1. 数据结构标准

数据集必须遵循 Robomimic 的层级结构：

dataset.hdf5
├── data/
│   ├── demo_0/
│   │   ├── obs/
│   │   │   ├── agentview_rgb   # (T, H, W, 3) uint8
│   │   │   └── qpos            # (T, D) float32
│   │   ├── actions             # (T, D) float32
│   │   └── language            # (Attribute) String 指令
│   └── ...

2. 数据转换工具

使用内置脚本将您的原始数据转换为标准 HDF5：

# 在项目根目录下运行
python -m roboimi.vla.scripts.convert_to_hdf5 \
  --input_dir /path/to/raw/images \
  --output_path ./data/demo.hdf5

3. 调试数据

如果不确定数据是否正确，使用可视化工具检查：

python -m roboimi.vla.scripts.visualize_data --dataset ./data/demo.hdf5

---

⚠️ 最佳实践 (Best Practices)

1. 绝对导入: 禁止使用 from . import xxx。请始终使用全路径 from roboimi.vla.models.backbones import SigLIPBackbone。
2. Hydra 插值: 在 agent/default.yaml 中，我们使用了 ${..embed_dim} 语法来确保所有子模块的维度一致。不要在子配置中硬编码维度数值。
3. HDF5 IO: 在 Dataset 类中，必须在 __getitem__ 内部打开 HDF5 文件。如果在 __init__ 中打开，多进程 DataLoader 会因无法序列化文件句柄而报错。
4. 接口导出:每当在 models/xxx/ 下添加新文件时，务必在对应的 __init__.py 中更新 __all__，以保持引用整洁。

---

Maintainer: VLA Framework Team