# iUMI Inference · 双臂UMI推理与部署

[English](README.md) · **简体中文**

Crobotia iUMI的模型推理与机器人部署代码，支持OpenPI（π0.5）和FLUX后端，包含配置模板、启动脚本及双臂名义URDF。

## 模型演示与 sample rollouts

本仓库提供 OpenPI π0.5 和 FLUX 的叠衣推理权重、部署源码、两段演示视频，以及每个后端各 1 条机器人执行样例。任务为 `fold_the_clothes`，硬件为 Flexiv Rizon 4 双臂与左右 ego 相机。

| 后端 | 演示视频（静音剪辑） | 完整 rollout | 帧数 / 30 Hz | 完整记录时长 |
| --- | --- | --- | ---: | ---: |
| OpenPI π0.5 | [4分15.4秒](demos/openpi_fold_demo.mp4) | [openpi_pi05](rollout_data/openpi_pi05) | 6,540 | 218 秒 |
| FLUX | [4分29.7秒](demos/flux_fold_demo.mp4) | [flux](rollout_data/flux) | 30,619 | 1,020.633 秒 |

**OpenPI π0.5**

<video src="https://huggingface.co/Crobotia/iUMI_Inference/resolve/main/demos/openpi_fold_demo.mp4" controls preload="metadata" width="100%"></video>

**FLUX**

<video src="https://huggingface.co/Crobotia/iUMI_Inference/resolve/main/demos/flux_fold_demo.mp4" controls preload="metadata" width="100%"></video>

视频用于快速浏览，剪辑时间轴不与 rollout 逐帧对应。两条 rollout 合计 37,159 帧、约 20分38.6秒，保留左右两路 960×960 / 30 Hz AV1 视频、16 维状态/动作、关节数据、双侧 IMU、标定和机器人 H5；不包含音频或 rear 相机。

只下载视频和 rollout，无需下载约 21.37 GB 的模型文件：

```bash
pip install -U huggingface_hub
hf download Crobotia/iUMI_Inference --local-dir iUMI_Inference --include "demos/*" "rollout_data/*"
```

媒体与 rollout 合计约 2.35 GB。分别以 `rollout_data/openpi_pi05/` 和 `rollout_data/flux/` 为独立 LeRobot v3 数据集根目录；两者都从 episode 0 开始。见 [rollout 读取说明](rollout_data/README.md) 和 [视频说明](demos/README.md)。这些文件存放在本模型仓库中，下载时使用默认 `repo_type="model"`。

## 模型用途与数据约定

两种后端用于双臂 UMI 衣物折叠的研究、离线接口集成和匹配设备上的部署验证。每次推理使用左右腕部图像、前一/当前位姿和任务文本，输出动作块，再由配套适配器完成坐标还原。

| 项目 | 约定 |
| --- | --- |
| 记录的状态 / 动作 | 16 维，左手在前；每手 `[x, y, z, qw, qx, qy, qz, gripper]`；米 / `wxyz` / 夹爪角度 |
| rollout 动作语义 | `action_mode=next_state`，`action[t] = observation.state[t+1]`（末帧除外），不是原始策略输出 |
| 模型内部 UMI 表示 | 两手合计 20 维；每手相对平移 3 + 旋转前两列 6 + 夹爪 1 |
| FLUX 动作块 | 32 步 / 30 Hz，使用双腕图像 |
| π0.5 动作块 | 50 个模型步；32 维模型头由 20 维 UMI 目标补齐，适配器还原为双臂目标位姿 |

记录位姿来源为 `robot_fk_cad`。相机光学坐标、机器人法兰坐标和模型相对表示需要显式转换，见 [输入输出与坐标约定](docs/INFERENCE_CONTRACT.md)。

## 训练来源、验证与局限

发布的 FLUX 为 step 14418 的 BF16 推理导出；π0.5 为 `pi05_umi_v3` 的 epoch 2 checkpoint。完整训练语料清单、划分和精确运行配置未随本次新增材料提供；这两条 rollout 是机器人执行样例，不能据此反推训练集。

本次新增材料通过文件哈希、视频流/包计数、Parquet 索引与下一帧动作关系、IMU 和 H5 结构检查，详见 [rollout 验证](docs/ROLLOUT_VALIDATION.md)。每个模型只有 1 条样例，未建立评测协议或任务成功率，也未运行完整运动质检。演示记录的确切 checkpoint 哈希、采样配置和剪辑到数据的逐帧映射未随材料提供，因此不能把当前启动脚本默认值视为视频的精确复现配方。

现有权重与源码的验证记录仍见 [验证范围](docs/VALIDATION.md)；此次上传没有重新运行 GPU 推理或机器人。实机部署需要匹配的编码器、厂商 SDK、设备标定、工作空间限制与现场验证。

## 获取代码

```bash
GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/Crobotia/iUMI_Inference
cd iUMI_Inference
```

| 目录 | 内容 |
| --- | --- |
| flexiv-umi-infer/ | 策略服务、机器人/相机/夹爪接口及数据记录 |
| flux-umi/ | FLUX模型与UMI数据表示 |
| openpi-umi/ | OpenPI模型、数据变换与客户端 |
| scripts/ | 模型服务和机器人端启动脚本 |
| configs/ | 推理配置模板 |
| checkpoints/ | FLUX BF16 和 π0.5 叠衣推理权重 |
| demos/ | 两段静音叠衣演示视频 |
| rollout_data/ | OpenPI π0.5 / FLUX 各 1 条完整执行记录 |

OpenPI使用Linux（glibc ≥2.31）；按`openpi-umi/.python-version`配置Python3.12，在该目录运行`uv sync --frozen --python 3.12`。

FLUX使用`scripts/run_flux_server.sh`，OpenPI使用`scripts/run_pi05_server.sh`，机器人端使用`scripts/run_robot.sh`。两类模型使用独立环境，按各子项目说明安装依赖。启动前配置模型资源路径及设备标定；机器人脚本默认为dry run。

运行前准备所需编码器、厂商SDK及设备标定。输入输出约定见`docs/INFERENCE_CONTRACT.md`，安装与设备配置见各子项目README。

## 服务部署关系

```mermaid
flowchart LR
  Robot["Flexiv 主机：robot_agent / 相机 / 本地伺服"] <-->|"WebSocket Observation / Action"| Server["GPU 主机：UMI policy_server"]
  Server --> Flux["FLUX 后端：flux-umi + FLUX checkpoint"]
  Server --> Pi["π0.5 后端：openpi-umi + OpenPI checkpoint"]
```

每个服务进程选择一个后端。`run_flux_server.sh` 和 `run_pi05_server.sh` 都启动 `umi_infer.policy_server`，分别调用对应的模型加载器；UMI 服务包装层位于 `flexiv-umi-infer/`。机器人端只运行 `robot_agent`，通过接口请求推理，不加载权重。模型输入编码、归一化和动作坐标还原在服务包装层执行；伺服执行在 Flexiv 主机上进行。

跨机器部署时，GPU 主机设置 `SERVER_HOST=0.0.0.0`、`SERVER_PORT=8765`；Flexiv 主机设置 `SERVER_URI=ws://<GPU主机地址>:8765`，再运行 `bash scripts/run_robot.sh`。同时部署两种模型时使用不同端口。设备标定路径 `UMI_INFER_DATA` 应在对应主机可访问；机器人脚本默认为 dry run。

请使用这套匹配的服务／客户端。OpenPI 原生 WebSocket 服务有不同的握手和数据封装，不能直接连接此 `robot_agent`。

## 下载权重与启动

两份权重已放在本仓库的 `checkpoints/` 下，共约 21.37 GB。在仓库根目录下载（需要 Hugging Face CLI）：

```bash
pip install -U huggingface_hub
hf download Crobotia/iUMI_Inference --local-dir . --include "checkpoints/*"
```

| 后端 | 启动脚本默认权重目录 | 模型文件大小 |
| --- | --- | --- |
| FLUX | `checkpoints/flux_fold_clothes_bf16/` | 13,894,466,592 B |
| π0.5 | `checkpoints/pi05_fold_clothes/epoch_2/` | 7,473,091,464 B |

启动脚本自动使用上述目录；也可用 `FLUX_EXPORT_DIR` 或 `PI05_CHECKPOINT_DIR` 指定其他位置。按后端安装依赖后，复制配置并填写资源、Python 环境和设备标定路径：

```bash
cp configs/inference.env.example .env
# 编辑 .env：FLUX_RESOURCES、TOKENIZER_DIR、UMI_INFER_DATA 及 Python 路径
set -a
source .env
set +a
bash scripts/run_flux_server.sh
# 或在 π0.5 环境中运行：bash scripts/run_pi05_server.sh
```

FLUX 仍需单独准备 `FLUX_RESOURCES/video_vae.safetensors` 和 `FLUX_RESOURCES/text_encoder/`；π0.5 仍需 `TOKENIZER_DIR/paligemma-3b-pt-224/tokenizer.model`。这些外部资源未包含在此次权重发布中。归一化参数随权重提供，启动时自动加载对应文件。

FLUX 导出未保存采样参数。当前脚本使用 Euler、4 步、shift 42、video/action guidance 1.0；这是可覆盖的脚本默认值，尚未验证叠衣效果。π0.5 默认 10 步去噪，`VLA_QUAT_PERM=0` 对应该权重的 `xyz_quat_wxyz` 配置。

权重已完成下载校验、文件结构及脱敏检查；CPU 检查不代表 GPU 推理或实机叠衣验证。详情见 [FLUX 权重说明](docs/CHECKPOINTS_FLUX.md)、[π0.5 权重说明](docs/CHECKPOINTS_PI05.md) 和 [验证范围](docs/VALIDATION.md)。

## 许可证

源码采用Apache-2.0许可，各子项目保留原有声明；Gemma、模型资源及厂商SDK遵循各自条款。详见[LICENSES.md](LICENSES.md)。

新增 demo 与 rollout 未附独立的媒体/数据许可声明；根目录 Apache-2.0 标签用于源码，不自动授予这些媒体、数据或模型资源的相同许可。使用或再分发前请向 Crobotia 确认适用条款；不将其他人工演示数据集的 CC BY-SA 声明套用于本次材料。