# VLA Framework: Vision-Language-Action Policy Framework

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

本框架基于 [Hydra](https://hydra.cc/) 进行配置管理，并采用 HDF5 作为标准数据格式。

---

## 🏗 架构概览 (Directory Structure)

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

```text
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. 环境依赖
请确保安装以下核心库：
```bash
pip install hydra-core h5py zarr diffusers transformers
```

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

```bash
# 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](https://robomimic.github.io/) 的层级结构：
```text
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：

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

### 3. 调试数据
如果不确定数据是否正确，使用可视化工具检查：
```bash
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*