diff --git a/.gitignore b/.gitignore index 6e9a55d..7f6fcf4 100644 --- a/.gitignore +++ b/.gitignore @@ -123,4 +123,9 @@ CLAUDE.md GEMINI.md # Copilot -.github/copilot-instructions.md \ No newline at end of file +.github/copilot-instructions.md + +.hydra/ + +# Local git worktrees +.worktrees/ diff --git a/README.en.md b/README.en.md deleted file mode 100644 index 024238a..0000000 --- a/README.en.md +++ /dev/null @@ -1,36 +0,0 @@ -# robo-imi-act - -#### Description -{**When you're done, you can delete the content in this README and update the file with details for others getting started with your repository**} - -#### Software Architecture -Software architecture description - -#### Installation - -1. xxxx -2. xxxx -3. xxxx - -#### Instructions - -1. xxxx -2. xxxx -3. xxxx - -#### Contribution - -1. Fork the repository -2. Create Feat_xxx branch -3. Commit your code -4. Create Pull Request - - -#### Gitee Feature - -1. You can use Readme\_XXX.md to support different languages, such as Readme\_en.md, Readme\_zh.md -2. Gitee blog [blog.gitee.com](https://blog.gitee.com) -3. Explore open source project [https://gitee.com/explore](https://gitee.com/explore) -4. The most valuable open source project [GVP](https://gitee.com/gvp) -5. The manual of Gitee [https://gitee.com/help](https://gitee.com/help) -6. The most popular members [https://gitee.com/gitee-stars/](https://gitee.com/gitee-stars/) diff --git a/README.md b/README.md index b72b112..b5d0744 100644 --- a/README.md +++ b/README.md @@ -1,39 +1,208 @@ -# robo-imi-act +# RoboIMI -#### 介绍 -{**以下是 Gitee 平台说明，您可以替换此简介** -Gitee 是 OSCHINA 推出的基于 Git 的代码托管平台（同时支持 SVN）。专为开发者提供稳定、高效、安全的云端软件开发协作平台 -无论是个人、团队、或是企业，都能够用 Gitee 实现代码托管、项目管理、协作开发。企业项目请看 [https://gitee.com/enterprises](https://gitee.com/enterprises)} +基于 MuJoCo 的机器人仿真与模仿学习框架，实现了使用扩散策略的视觉-语言-动作（VLA）模型，用于机器人操作任务。 -#### 软件架构 -软件架构说明 +## 主要特性 +- **多机器人平台支持**：支持 Diana 和 vx300s 机械臂，可扩展至其他机器人 +- **扩散策略**：采用最先进的扩散模型（DDPM/DDIM）进行动作序列预测 +- **视觉-语言-动作模型**：使用 ResNet-18 视觉骨干网络和空间 softmax 进行视觉特征提取 +- **灵活的控制模式**：支持关节空间和末端执行器（笛卡尔）控制 +- **Hydra 配置系统**：模块化配置系统，便于实验 +- **HDF5 数据集格式**：高效存储和加载演示数据 +- **单臂和双臂任务**：支持单臂和双臂操作任务 -#### 安装教程 +## 安装 -1. xxxx -2. xxxx -3. xxxx +### 环境要求 -#### 使用说明 +- Python 3.8+ +- 支持 CUDA 的 GPU（训练时推荐） +- Conda 或 Miniconda -1. xxxx -2. xxxx -3. xxxx +### 安装步骤 -#### 参与贡献 +```bash +# 克隆仓库 +git clone +cd robo-imi-act -1. Fork 本仓库 -2. 新建 Feat_xxx 分支 -3. 提交代码 -4. 新建 Pull Request +# 创建并激活 conda 环境 +conda env create -f environment.yml +conda activate roboimi +# 以开发模式安装包 +pip install -e . +``` -#### 特技 +## 快速开始 -1. 使用 Readme\_XXX.md 来支持不同的语言，例如 Readme\_en.md, Readme\_zh.md -2. Gitee 官方博客 [blog.gitee.com](https://blog.gitee.com) -3. 你可以 [https://gitee.com/explore](https://gitee.com/explore) 这个地址来了解 Gitee 上的优秀开源项目 -4. [GVP](https://gitee.com/gvp) 全称是 Gitee 最有价值开源项目，是综合评定出的优秀开源项目 -5. Gitee 官方提供的使用手册 [https://gitee.com/help](https://gitee.com/help) -6. Gitee 封面人物是一档用来展示 Gitee 会员风采的栏目 [https://gitee.com/gitee-stars/](https://gitee.com/gitee-stars/) +### 1. 数据采集 + +在仿真环境中记录演示轨迹： + +```bash +# 为 vx300s 机器人记录轨迹 +python roboimi/demos/record_sim_episodes.py + +# 为 Diana 机器人记录轨迹 +python roboimi/demos/diana_record_sim_episodes.py +``` + +轨迹数据以 HDF5 文件格式保存，包含机器人状态、动作和相机观测。 + +### 2. 计算数据集统计信息 + +训练前需要计算归一化统计数据： + +```bash +python roboimi/vla/scripts/calculate_stats.py +``` + +该命令会生成 `data_stats.pkl` 文件，包含动作和观测的均值/标准差或最小值/最大值。 + +### 3. 训练 VLA 模型 + +使用采集的数据训练视觉-语言-动作模型： + +```bash +# 使用默认配置训练 +python roboimi/demos/vla_scripts/train_vla.py + +# 覆盖特定参数 +python roboimi/demos/vla_scripts/train_vla.py train.batch_size=32 train.lr=5e-5 train.max_steps=50000 + +# 使用不同的模型架构 +python roboimi/demos/vla_scripts/train_vla.py agent=resnet_diffusion data=resnet_dataset +``` + +训练输出保存至 `outputs/<日期>/<时间>/`，模型检查点保存至 `checkpoints/`。 + +### 4. 评估模型 + +在仿真环境中评估训练好的模型： + +```bash +# 使用默认配置评估（使用最佳检查点） +python roboimi/demos/vla_scripts/eval_vla.py + +# 指定检查点和评估轮数 +python roboimi/demos/vla_scripts/eval_vla.py eval.ckpt_path=checkpoints/vla_model_step_8000.pt eval.num_episodes=5 + +# 启用动作平滑以获得更流畅的执行 +python roboimi/demos/vla_scripts/eval_vla.py eval.use_smoothing=true eval.smooth_alpha=0.5 +``` + +## 项目结构 + +``` +robo-imi-act/ +├── roboimi/ +│ ├── assets/ # 机器人模型和资源 +│ │ ├── models/manipulators/ # URDF 和 MuJoCo XML 文件 +│ │ └── robots/ # 机器人抽象类 +│ ├── envs/ # 仿真环境 +│ │ ├── mujoco_base.py # MuJoCo 环境基类 +│ │ ├── single_base.py # 单臂任务基类 +│ │ └── double_base.py # 双臂任务基类 +│ ├── vla/ # 视觉-语言-动作模型 +│ │ ├── agent.py # VLAAgent（训练与推理） +│ │ ├── models/ +│ │ │ ├── backbones/ # 视觉编码器（ResNet 等） +│ │ │ └── heads/ # 策略头（扩散 UNet1D） +│ │ ├── conf/ # Hydra 配置文件 +│ │ └── scripts/ # 训练和工具脚本 +│ └── demos/ # 演示脚本和示例 +├── checkpoints/ # 保存的模型检查点 +├── outputs/ # 训练输出（Hydra） +├── environment.yml # Conda 环境配置 +└── CLAUDE.md # Claude Code 开发指南 +``` + +## 架构设计 + +### VLA 训练流程 + +``` +HDF5 轨迹数据 → Dataset → DataLoader → VLAAgent → 模型检查点 +``` + +**模型组件**： +- **视觉骨干网络**：ResNet-18 + 空间 softmax，用于从相机图像中提取视觉特征 +- **扩散头**：条件 UNet1D，使用 DDPM/DDIM 预测动作序列 +- **VLAAgent**：组合视觉编码器和扩散策略，处理训练和推理 + +### 配置系统 + +基于 Hydra 的配置文件位于 `roboimi/vla/conf/`： +- `config.yaml`：主要训练配置（批次大小、学习率、设备） +- `agent/resnet_diffusion.yaml`：模型架构（动作维度、观测维度、时间窗口） +- `data/resnet_dataset.yaml`：数据集路径、相机名称、归一化类型 +- `eval/eval.yaml`：评估设置（检查点路径、轮数、平滑参数） + +使用配置插值保持一致性：`${agent.obs_horizon}` + +### 数据集格式 + +HDF5 轨迹文件（`episode_*.hdf5`）包含： +- `action`：机器人动作 `[T, action_dim]` +- `observations/qpos`：关节位置 `[T, obs_dim]` +- `observations/images/`：相机图像 `[T, H, W, C]` + +统计文件（`data_stats.pkl`）存储归一化参数（最小值/最大值/均值/标准差）。 + +## 开发指南 + +### 添加新机器人 + +1. 在 `roboimi/assets/models/manipulators//` 创建 URDF/XML 文件 +2. 在 `roboimi/assets/robots/.py` 定义机器人类（继承自 `arm_base.py`） +3. 在 `roboimi/envs/_*.py` 创建环境类 +4. 如需要，在常量中注册机器人 + +### 修改 VLA 架构 + +1. **自定义骨干网络**：在 `roboimi/vla/models/backbones/` 创建新类，继承 `VLABackbone` +2. **自定义头部**：在 `roboimi/vla/models/heads/` 创建新类，继承 `VLAHead` +3. **更新配置**：在 `roboimi/vla/conf/agent/` 添加新的 YAML 文件 +4. **接口定义**：参考 `roboimi/vla/core/interfaces.py` 的抽象基类 + +### 训练最佳实践 + +- 采集新数据后务必运行 `calculate_stats.py` +- 训练时会归一化输入/输出；推理时使用检查点中保存的统计信息进行反归一化 +- 模型预测 `pred_horizon` 步，但只执行前 `action_horizon` 步 +- 推理使用 DDIM（10 步）快速采样；训练使用 DDPM（100 步） +- 监控验证损失以防止过拟合 + +## 技术细节 + +- **坐标系**：关节空间（qpos）或末端执行器空间（xyz + rpy + 夹爪） +- **动作时间窗口**：`obs_horizon` 为观测窗口，`pred_horizon` 为预测窗口，`action_horizon` 为执行窗口 +- **归一化**：对稳定训练至关重要 - 训练前务必计算统计信息 +- **推理加速**：使用 DDIM 调度器，比训练时的 DDPM 快 10 倍 +- **设备配置**：通过 `train.device` 配置（cuda/cpu） + +## 许可证 + +[在此添加许可证信息] + +## 引用 + +如果您在研究中使用了本代码库，请引用： + +```bibtex +[在此添加引用信息] +``` + +## 贡献 + +欢迎贡献！请随时提交 Pull Request 或开启 Issue。 + +## 致谢 + +本项目基于以下开源项目构建： +- [MuJoCo](https://mujoco.org/) - 物理仿真引擎 +- [PyTorch](https://pytorch.org/) - 深度学习框架 +- [Hydra](https://hydra.cc/) - 配置管理系统 +- [Diffusers](https://github.com/huggingface/diffusers) - 扩散模型库 diff --git a/check_all_episodes.py b/check_all_episodes.py new file mode 100644 index 0000000..2734216 --- /dev/null +++ b/check_all_episodes.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +""" +检查所有 episode 的重复帧情况 + +找出哪些 episode 有问题，需要删除或重新收集 +""" +import os +import h5py +import glob +import numpy as np + + +def check_all_episodes(): + """检查所有 episode 的质量""" + + dataset_dir = "roboimi/demos/dataset/sim_transfer" + episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) + episode_files = sorted(episode_files, key=lambda x: int(x.split('_')[-1].replace('.hdf5', ''))) + + print("="*80) + print("所有 Episode 质量检查") + print("="*80) + + good_episodes = [] + bad_episodes = [] + + for ep_idx, ep_file in enumerate(episode_files): + ep_name = os.path.basename(ep_file).replace('.hdf5', '') + + try: + with h5py.File(ep_file, 'r') as f: + img_path = '/observations/images/top' + if img_path not in f: + continue + + images = f[img_path][:] + + # 检查前 50 帧的重复情况 + check_frames = min(50, len(images)) + duplicate_count = 0 + + for i in range(check_frames - 1): + img1 = images[i] + img2 = images[i + 1] + diff = np.mean(np.abs(img1.astype(float) - img2.astype(float))) + + if diff < 1.0: # 重复 + duplicate_count += 1 + + duplicate_rate = duplicate_count / check_frames * 100 + + # 判断质量 + if duplicate_rate > 10: # 超过10%重复 + bad_episodes.append((ep_idx, ep_name, duplicate_rate, duplicate_count)) + status = "❌" + else: + good_episodes.append((ep_idx, ep_name, duplicate_rate, duplicate_count)) + status = "✅" + + print(f"{status} Episode {ep_idx:2d}: {duplicate_rate:5.1f}% 重复 ({duplicate_count:2d}/{check_frames}) - {ep_name}") + + except Exception as e: + print(f"❌ Episode {ep_idx}: 错误 - {e}") + + # 总结 + print("\n" + "="*80) + print("总结") + print("="*80) + print(f"总共检查: {len(episode_files)} 个 episodes") + print(f"正常的: {len(good_episodes)} 个 ✅") + print(f"有问题的: {len(bad_episodes)} 个 ❌") + + if bad_episodes: + print(f"\n有问题的 episodes:") + for ep_idx, ep_name, rate, count in bad_episodes: + print(f" - episode_{ep_idx}.hdf5: {rate:.1f}% 重复") + + print(f"\n删除命令:") + ep_names = [name for _, name, _, _ in bad_episodes] + print(f" rm " + " ".join([f"{dataset_dir}/{name}.hdf5" for name in ep_names])) + + print(f"\n建议:") + if len(bad_episodes) > 0: + print(f" 1. 删除有问题的 {len(bad_episodes)} 个 episodes") + print(f" 2. 重新收集数据，或使用剩余的 {len(good_episodes)} 个正常 episodes") + else: + print(f" ✅ 所有 episodes 都正常，可以直接使用！") + + +if __name__ == "__main__": + check_all_episodes() diff --git a/check_specific_frames.py b/check_specific_frames.py new file mode 100644 index 0000000..ce93d35 --- /dev/null +++ b/check_specific_frames.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +""" +检查特定帧的图像 - 用于验证数据记录问题 + +功能： +1. 提取每个 episode 的第 0、1、2 帧图像 +2. 对比不同 episode 的相同帧号 +3. 保存图像供人工检查 +""" +import os +import h5py +import glob +import cv2 +import numpy as np + + +def check_specific_frames(frame_indices=[0, 1, 2], camera='top', num_episodes=10): + """ + 检查特定帧的图像和 qpos + + Args: + frame_indices: 要检查的帧索引列表 + camera: 相机名称 + num_episodes: 要检查的 episode 数量 + """ + + dataset_dir = "roboimi/demos/dataset/sim_transfer" + episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) + # 按数字排序 + episode_files = sorted(episode_files, key=lambda x: int(x.split('_')[-1].replace('.hdf5', ''))) + + # 创建输出目录 + output_dir = f'/tmp/dataset_frames' + os.makedirs(output_dir, exist_ok=True) + + print(f"检查前 {min(num_episodes, len(episode_files))} 个 episode 的特定帧") + print(f"帧索引: {frame_indices}") + print(f"相机: {camera}") + print("="*80) + + # 收集所有数据 + for ep_idx in range(min(num_episodes, len(episode_files))): + ep_file = episode_files[ep_idx] + ep_name = os.path.basename(ep_file).replace('.hdf5', '') + + try: + with h5py.File(ep_file, 'r') as f: + # 读取 qpos + qpos = f['/observations/qpos'][:] + + # 读取图像 + img_path = f'/observations/images/{camera}' + if img_path not in f: + print(f"Episode {ep_name}: 相机 {camera} 不存在") + continue + + images = f[img_path][:] + + print(f"\nEpisode {ep_name}:") + print(f" 总帧数: {len(images)}") + + # 保存指定帧 + for frame_idx in frame_indices: + if frame_idx >= len(images): + print(f" 帧 {frame_idx}: 超出范围") + continue + + # 保存图像 + img = images[frame_idx] + filename = f"{output_dir}/ep{ep_idx:02d}_frame{frame_idx:03d}.png" + cv2.imwrite(filename, img) + + # 打印 qpos + q = qpos[frame_idx] + print(f" 帧 {frame_idx}: qpos[0:3]=[{q[0]:6.2f}, {q[1]:6.2f}, {q[2]:6.2f}], qpos[3]={q[3]:6.2f} → {filename}") + + except Exception as e: + print(f"Episode {ep_name}: 错误 - {e}") + + print("\n" + "="*80) + print(f"✅ 所有图像已保存到: {output_dir}") + print(f"\n查看方法:") + print(f" eog {output_dir}/*.png") + print(f" ") + print(f" # 或对比特定帧:") + print(f" eog {output_dir}/*_frame000.png # 所有 episode 的第 0 帧") + print(f" eog {output_dir}/*_frame001.png # 所有 episode 的第 1 帧") + print(f" eog {output_dir}/*_frame002.png # 所有 episode 的第 2 帧") + + +def compare_frame_across_episodes(frame_idx=0, camera='top', num_episodes=10): + """ + 并排对比所有 episode 的某一帧 + + 生成一个大的对比图，包含所有 episode 的指定帧 + """ + + dataset_dir = "roboimi/demos/dataset/sim_transfer" + episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) + episode_files = sorted(episode_files, key=lambda x: int(x.split('_')[-1].replace('.hdf5', ''))) + + num_compare = min(num_episodes, len(episode_files)) + cols = 5 # 每行 5 个 + rows = (num_compare + cols - 1) // cols + + # 创建输出目录 + output_dir = f'/tmp/dataset_frames' + os.makedirs(output_dir, exist_ok=True) + + print(f"生成对比图: 所有 Episode 的第 {frame_idx} 帧") + print("="*80) + + # 收集图像 + images_compare = [] + qpos_list = [] + + for ep_idx in range(num_compare): + ep_file = episode_files[ep_idx] + ep_name = os.path.basename(ep_file).replace('.hdf5', '') + + try: + with h5py.File(ep_file, 'r') as f: + qpos = f['/observations/qpos'][:] + img_path = f'/observations/images/{camera}' + + if img_path in f and frame_idx < f[img_path].shape[0]: + img = f[img_path][frame_idx] + images_compare.append(img) + qpos_list.append(qpos[frame_idx]) + print(f"Episode {ep_name}: qpos[0:3]=[{qpos[frame_idx][0]:.2f}, {qpos[frame_idx][1]:.2f}, {qpos[frame_idx][2]:.2f}]") + + except Exception as e: + print(f"Episode {ep_name}: 错误 - {e}") + + if not images_compare: + print("❌ 没有收集到图像") + return + + # 获取图像尺寸 + h, w = images_compare[0].shape[:2] + + # 创建对比图 + compare_img = np.zeros((rows * h + 50, cols * w, 3), dtype=np.uint8) + + for i, (img, qpos) in enumerate(zip(images_compare, qpos_list)): + row = i // cols + col = i % cols + + y_start = row * h + 30 + y_end = y_start + h + x_start = col * w + x_end = x_start + w + + # 调整大小（如果需要） + if img.shape[:2] != (h, w): + img = cv2.resize(img, (w, h)) + + compare_img[y_start:y_end, x_start:x_end] = img + + # 添加信息 + ep_name = f"Ep {i}" + cv2.putText(compare_img, ep_name, (x_start + 10, row * h + 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2) + cv2.putText(compare_img, f"qpos[3]={qpos[3]:.2f}", (x_start + 10, y_end - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + + # 保存对比图 + output_path = f"{output_dir}/compare_frame{frame_idx:03d}.png" + cv2.imwrite(output_path, compare_img) + + print(f"\n✅ 对比图已保存: {output_path}") + print(f" 查看方法: eog {output_path}") + + +if __name__ == "__main__": + import sys + + print("="*80) + print("特定帧检查工具") + print("="*80) + + if len(sys.argv) > 1: + frame_idx = int(sys.argv[1]) + compare_frame_across_episodes(frame_idx=frame_idx, camera='top', num_episodes=10) + else: + # 默认检查第 0、1、2 帧 + check_specific_frames(frame_indices=[0, 1, 2], camera='top', num_episodes=10) + + print("\n" + "="*80) + print("生成对比图...") + print("="*80) + + # 生成第 0 帧的对比图 + compare_frame_across_episodes(frame_idx=0, camera='top', num_episodes=10) + compare_frame_across_episodes(frame_idx=1, camera='top', num_episodes=10) + compare_frame_across_episodes(frame_idx=2, camera='top', num_episodes=10) + + print("\n" + "="*80) + print("其他用法:") + print(" python check_specific_frames.py 0 # 只检查第 0 帧") + print(" python check_specific_frames.py 1 # 只检查第 1 帧") + print("="*80) diff --git a/diffusion/configuration_diffusion.py b/diffusion/configuration_diffusion.py new file mode 100644 index 0000000..5456943 --- /dev/null +++ b/diffusion/configuration_diffusion.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python + +# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab, +# and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from dataclasses import dataclass, field + +from lerobot.configs.policies import PreTrainedConfig +from lerobot.configs.types import NormalizationMode +from lerobot.optim.optimizers import AdamConfig +from lerobot.optim.schedulers import DiffuserSchedulerConfig + + +@PreTrainedConfig.register_subclass("diffusion") +@dataclass +class DiffusionConfig(PreTrainedConfig): + """Configuration class for DiffusionPolicy. + + Defaults are configured for training with PushT providing proprioceptive and single camera observations. + + The parameters you will most likely need to change are the ones which depend on the environment / sensors. + Those are: `input_shapes` and `output_shapes`. + + Notes on the inputs and outputs: + - "observation.state" is required as an input key. + - Either: + - At least one key starting with "observation.image is required as an input. + AND/OR + - The key "observation.environment_state" is required as input. + - If there are multiple keys beginning with "observation.image" they are treated as multiple camera + views. Right now we only support all images having the same shape. + - "action" is required as an output key. + + Args: + n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the + current step and additional steps going back). + horizon: Diffusion model action prediction size as detailed in `DiffusionPolicy.select_action`. + n_action_steps: The number of action steps to run in the environment for one invocation of the policy. + See `DiffusionPolicy.select_action` for more details. + input_shapes: A dictionary defining the shapes of the input data for the policy. The key represents + the input data name, and the value is a list indicating the dimensions of the corresponding data. + For example, "observation.image" refers to an input from a camera with dimensions [3, 96, 96], + indicating it has three color channels and 96x96 resolution. Importantly, `input_shapes` doesn't + include batch dimension or temporal dimension. + output_shapes: A dictionary defining the shapes of the output data for the policy. The key represents + the output data name, and the value is a list indicating the dimensions of the corresponding data. + For example, "action" refers to an output shape of [14], indicating 14-dimensional actions. + Importantly, `output_shapes` doesn't include batch dimension or temporal dimension. + input_normalization_modes: A dictionary with key representing the modality (e.g. "observation.state"), + and the value specifies the normalization mode to apply. The two available modes are "mean_std" + which subtracts the mean and divides by the standard deviation and "min_max" which rescale in a + [-1, 1] range. + output_normalization_modes: Similar dictionary as `normalize_input_modes`, but to unnormalize to the + original scale. Note that this is also used for normalizing the training targets. + vision_backbone: Name of the torchvision resnet backbone to use for encoding images. + crop_shape: (H, W) shape to crop images to as a preprocessing step for the vision backbone. Must fit + within the image size. If None, no cropping is done. + crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval + mode). + pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone. + `None` means no pretrained weights. + use_group_norm: Whether to replace batch normalization with group normalization in the backbone. + The group sizes are set to be about 16 (to be precise, feature_dim // 16). + spatial_softmax_num_keypoints: Number of keypoints for SpatialSoftmax. + use_separate_rgb_encoders_per_camera: Whether to use a separate RGB encoder for each camera view. + down_dims: Feature dimension for each stage of temporal downsampling in the diffusion modeling Unet. + You may provide a variable number of dimensions, therefore also controlling the degree of + downsampling. + kernel_size: The convolutional kernel size of the diffusion modeling Unet. + n_groups: Number of groups used in the group norm of the Unet's convolutional blocks. + diffusion_step_embed_dim: The Unet is conditioned on the diffusion timestep via a small non-linear + network. This is the output dimension of that network, i.e., the embedding dimension. + use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning. + Bias modulation is used be default, while this parameter indicates whether to also use scale + modulation. + noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"]. + num_train_timesteps: Number of diffusion steps for the forward diffusion schedule. + beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers. + beta_start: Beta value for the first forward-diffusion step. + beta_end: Beta value for the last forward-diffusion step. + prediction_type: The type of prediction that the diffusion modeling Unet makes. Choose from "epsilon" + or "sample". These have equivalent outcomes from a latent variable modeling perspective, but + "epsilon" has been shown to work better in many deep neural network settings. + clip_sample: Whether to clip the sample to [-`clip_sample_range`, +`clip_sample_range`] for each + denoising step at inference time. WARNING: you will need to make sure your action-space is + normalized to fit within this range. + clip_sample_range: The magnitude of the clipping range as described above. + num_inference_steps: Number of reverse diffusion steps to use at inference time (steps are evenly + spaced). If not provided, this defaults to be the same as `num_train_timesteps`. + do_mask_loss_for_padding: Whether to mask the loss when there are copy-padded actions. See + `LeRobotDataset` and `load_previous_and_future_frames` for more information. Note, this defaults + to False as the original Diffusion Policy implementation does the same. + """ + + # Inputs / output structure. + n_obs_steps: int = 2 + horizon: int = 16 + n_action_steps: int = 8 + + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.MEAN_STD, + "STATE": NormalizationMode.MIN_MAX, + "ACTION": NormalizationMode.MIN_MAX, + } + ) + + # The original implementation doesn't sample frames for the last 7 steps, + # which avoids excessive padding and leads to improved training results. + drop_n_last_frames: int = 7 # horizon - n_action_steps - n_obs_steps + 1 + + # Architecture / modeling. + # Vision backbone. + vision_backbone: str = "resnet18" + crop_shape: tuple[int, int] | None = (84, 84) + crop_is_random: bool = True + pretrained_backbone_weights: str | None = None + use_group_norm: bool = True + spatial_softmax_num_keypoints: int = 32 + use_separate_rgb_encoder_per_camera: bool = False + # Unet. + down_dims: tuple[int, ...] = (512, 1024, 2048) + kernel_size: int = 5 + n_groups: int = 8 + diffusion_step_embed_dim: int = 128 + use_film_scale_modulation: bool = True + # Noise scheduler. + noise_scheduler_type: str = "DDPM" + num_train_timesteps: int = 100 + beta_schedule: str = "squaredcos_cap_v2" + beta_start: float = 0.0001 + beta_end: float = 0.02 + prediction_type: str = "epsilon" + clip_sample: bool = True + clip_sample_range: float = 1.0 + + # Inference + num_inference_steps: int | None = None + + # Loss computation + do_mask_loss_for_padding: bool = False + + # Training presets + optimizer_lr: float = 1e-4 + optimizer_betas: tuple = (0.95, 0.999) + optimizer_eps: float = 1e-8 + optimizer_weight_decay: float = 1e-6 + scheduler_name: str = "cosine" + scheduler_warmup_steps: int = 500 + + def __post_init__(self): + super().__post_init__() + + """Input validation (not exhaustive).""" + if not self.vision_backbone.startswith("resnet"): + raise ValueError( + f"`vision_backbone` must be one of the ResNet variants. Got {self.vision_backbone}." + ) + + supported_prediction_types = ["epsilon", "sample"] + if self.prediction_type not in supported_prediction_types: + raise ValueError( + f"`prediction_type` must be one of {supported_prediction_types}. Got {self.prediction_type}." + ) + supported_noise_schedulers = ["DDPM", "DDIM"] + if self.noise_scheduler_type not in supported_noise_schedulers: + raise ValueError( + f"`noise_scheduler_type` must be one of {supported_noise_schedulers}. " + f"Got {self.noise_scheduler_type}." + ) + + # Check that the horizon size and U-Net downsampling is compatible. + # U-Net downsamples by 2 with each stage. + downsampling_factor = 2 ** len(self.down_dims) + if self.horizon % downsampling_factor != 0: + raise ValueError( + "The horizon should be an integer multiple of the downsampling factor (which is determined " + f"by `len(down_dims)`). Got {self.horizon=} and {self.down_dims=}" + ) + + def get_optimizer_preset(self) -> AdamConfig: + return AdamConfig( + lr=self.optimizer_lr, + betas=self.optimizer_betas, + eps=self.optimizer_eps, + weight_decay=self.optimizer_weight_decay, + ) + + def get_scheduler_preset(self) -> DiffuserSchedulerConfig: + return DiffuserSchedulerConfig( + name=self.scheduler_name, + num_warmup_steps=self.scheduler_warmup_steps, + ) + + def validate_features(self) -> None: + if len(self.image_features) == 0 and self.env_state_feature is None: + raise ValueError("You must provide at least one image or the environment state among the inputs.") + + if self.crop_shape is not None: + for key, image_ft in self.image_features.items(): + if self.crop_shape[0] > image_ft.shape[1] or self.crop_shape[1] > image_ft.shape[2]: + raise ValueError( + f"`crop_shape` should fit within the images shapes. Got {self.crop_shape} " + f"for `crop_shape` and {image_ft.shape} for " + f"`{key}`." + ) + + # Check that all input images have the same shape. + if len(self.image_features) > 0: + first_image_key, first_image_ft = next(iter(self.image_features.items())) + for key, image_ft in self.image_features.items(): + if image_ft.shape != first_image_ft.shape: + raise ValueError( + f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match." + ) + + @property + def observation_delta_indices(self) -> list: + return list(range(1 - self.n_obs_steps, 1)) + + @property + def action_delta_indices(self) -> list: + return list(range(1 - self.n_obs_steps, 1 - self.n_obs_steps + self.horizon)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/diffusion/modeling_diffusion.py b/diffusion/modeling_diffusion.py new file mode 100644 index 0000000..1fdc76f --- /dev/null +++ b/diffusion/modeling_diffusion.py @@ -0,0 +1,764 @@ +#!/usr/bin/env python + +# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab, +# and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Diffusion Policy as per "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion" + +TODO(alexander-soare): + - Remove reliance on diffusers for DDPMScheduler and LR scheduler. +""" + +import math +from collections import deque +from collections.abc import Callable + +import einops +import numpy as np +import torch +import torch.nn.functional as F # noqa: N812 +import torchvision +from diffusers.schedulers.scheduling_ddim import DDIMScheduler +from diffusers.schedulers.scheduling_ddpm import DDPMScheduler +from torch import Tensor, nn + +from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.policies.utils import ( + get_device_from_parameters, + get_dtype_from_parameters, + get_output_shape, + populate_queues, +) +from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE + + +class DiffusionPolicy(PreTrainedPolicy): + """ + Diffusion Policy as per "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion" + (paper: https://huggingface.co/papers/2303.04137, code: https://github.com/real-stanford/diffusion_policy). + """ + + config_class = DiffusionConfig + name = "diffusion" + + def __init__( + self, + config: DiffusionConfig, + **kwargs, + ): + """ + Args: + config: Policy configuration class instance or None, in which case the default instantiation of + the configuration class is used. + dataset_stats: Dataset statistics to be used for normalization. If not passed here, it is expected + that they will be passed with a call to `load_state_dict` before the policy is used. + """ + super().__init__(config) + config.validate_features() + self.config = config + + # queues are populated during rollout of the policy, they contain the n latest observations and actions + self._queues = None + + self.diffusion = DiffusionModel(config) + + self.reset() + + def get_optim_params(self) -> dict: + return self.diffusion.parameters() + + def reset(self): + """Clear observation and action queues. Should be called on `env.reset()`""" + self._queues = { + OBS_STATE: deque(maxlen=self.config.n_obs_steps), + ACTION: deque(maxlen=self.config.n_action_steps), + } + if self.config.image_features: + self._queues[OBS_IMAGES] = deque(maxlen=self.config.n_obs_steps) + if self.config.env_state_feature: + self._queues[OBS_ENV_STATE] = deque(maxlen=self.config.n_obs_steps) + + @torch.no_grad() + def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: + """Predict a chunk of actions given environment observations.""" + # stack n latest observations from the queue + batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} + actions = self.diffusion.generate_actions(batch, noise=noise) + + return actions + + @torch.no_grad() + def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: + """Select a single action given environment observations. + + This method handles caching a history of observations and an action trajectory generated by the + underlying diffusion model. Here's how it works: + - `n_obs_steps` steps worth of observations are cached (for the first steps, the observation is + copied `n_obs_steps` times to fill the cache). + - The diffusion model generates `horizon` steps worth of actions. + - `n_action_steps` worth of actions are actually kept for execution, starting from the current step. + Schematically this looks like: + ---------------------------------------------------------------------------------------------- + (legend: o = n_obs_steps, h = horizon, a = n_action_steps) + |timestep | n-o+1 | n-o+2 | ..... | n | ..... | n+a-1 | n+a | ..... | n-o+h | + |observation is used | YES | YES | YES | YES | NO | NO | NO | NO | NO | + |action is generated | YES | YES | YES | YES | YES | YES | YES | YES | YES | + |action is used | NO | NO | NO | YES | YES | YES | NO | NO | NO | + ---------------------------------------------------------------------------------------------- + Note that this means we require: `n_action_steps <= horizon - n_obs_steps + 1`. Also, note that + "horizon" may not the best name to describe what the variable actually means, because this period is + actually measured from the first observation which (if `n_obs_steps` > 1) happened in the past. + """ + # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out + if ACTION in batch: + batch.pop(ACTION) + + if self.config.image_features: + batch = dict(batch) # shallow copy so that adding a key doesn't modify the original + batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) + # NOTE: It's important that this happens after stacking the images into a single key. + self._queues = populate_queues(self._queues, batch) + + if len(self._queues[ACTION]) == 0: + actions = self.predict_action_chunk(batch, noise=noise) + self._queues[ACTION].extend(actions.transpose(0, 1)) + + action = self._queues[ACTION].popleft() + return action + + def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]: + """Run the batch through the model and compute the loss for training or validation.""" + if self.config.image_features: + batch = dict(batch) # shallow copy so that adding a key doesn't modify the original + batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) + loss = self.diffusion.compute_loss(batch) + # no output_dict so returning None + return loss, None + + +def _make_noise_scheduler(name: str, **kwargs: dict) -> DDPMScheduler | DDIMScheduler: + """ + Factory for noise scheduler instances of the requested type. All kwargs are passed + to the scheduler. + """ + if name == "DDPM": + return DDPMScheduler(**kwargs) + elif name == "DDIM": + return DDIMScheduler(**kwargs) + else: + raise ValueError(f"Unsupported noise scheduler type {name}") + + +class DiffusionModel(nn.Module): + def __init__(self, config: DiffusionConfig): + super().__init__() + self.config = config + + # Build observation encoders (depending on which observations are provided). + global_cond_dim = self.config.robot_state_feature.shape[0] + if self.config.image_features: + num_images = len(self.config.image_features) + if self.config.use_separate_rgb_encoder_per_camera: + encoders = [DiffusionRgbEncoder(config) for _ in range(num_images)] + self.rgb_encoder = nn.ModuleList(encoders) + global_cond_dim += encoders[0].feature_dim * num_images + else: + self.rgb_encoder = DiffusionRgbEncoder(config) + global_cond_dim += self.rgb_encoder.feature_dim * num_images + if self.config.env_state_feature: + global_cond_dim += self.config.env_state_feature.shape[0] + + self.unet = DiffusionConditionalUnet1d(config, global_cond_dim=global_cond_dim * config.n_obs_steps) + + self.noise_scheduler = _make_noise_scheduler( + config.noise_scheduler_type, + num_train_timesteps=config.num_train_timesteps, + beta_start=config.beta_start, + beta_end=config.beta_end, + beta_schedule=config.beta_schedule, + clip_sample=config.clip_sample, + clip_sample_range=config.clip_sample_range, + prediction_type=config.prediction_type, + ) + + if config.num_inference_steps is None: + self.num_inference_steps = self.noise_scheduler.config.num_train_timesteps + else: + self.num_inference_steps = config.num_inference_steps + + # ========= inference ============ + def conditional_sample( + self, + batch_size: int, + global_cond: Tensor | None = None, + generator: torch.Generator | None = None, + noise: Tensor | None = None, + ) -> Tensor: + device = get_device_from_parameters(self) + dtype = get_dtype_from_parameters(self) + + # Sample prior. + sample = ( + noise + if noise is not None + else torch.randn( + size=(batch_size, self.config.horizon, self.config.action_feature.shape[0]), + dtype=dtype, + device=device, + generator=generator, + ) + ) + + self.noise_scheduler.set_timesteps(self.num_inference_steps) + + for t in self.noise_scheduler.timesteps: + # Predict model output. + model_output = self.unet( + sample, + torch.full(sample.shape[:1], t, dtype=torch.long, device=sample.device), + global_cond=global_cond, + ) + # Compute previous image: x_t -> x_t-1 + sample = self.noise_scheduler.step(model_output, t, sample, generator=generator).prev_sample + + return sample + + def _prepare_global_conditioning(self, batch: dict[str, Tensor]) -> Tensor: + """Encode image features and concatenate them all together along with the state vector.""" + batch_size, n_obs_steps = batch[OBS_STATE].shape[:2] + global_cond_feats = [batch[OBS_STATE]] + # Extract image features. + if self.config.image_features: + if self.config.use_separate_rgb_encoder_per_camera: + # Combine batch and sequence dims while rearranging to make the camera index dimension first. + images_per_camera = einops.rearrange(batch[OBS_IMAGES], "b s n ... -> n (b s) ...") + img_features_list = torch.cat( + [ + encoder(images) + for encoder, images in zip(self.rgb_encoder, images_per_camera, strict=True) + ] + ) + # Separate batch and sequence dims back out. The camera index dim gets absorbed into the + # feature dim (effectively concatenating the camera features). + img_features = einops.rearrange( + img_features_list, "(n b s) ... -> b s (n ...)", b=batch_size, s=n_obs_steps + ) + else: + # Combine batch, sequence, and "which camera" dims before passing to shared encoder. + img_features = self.rgb_encoder( + einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ...") + ) + # Separate batch dim and sequence dim back out. The camera index dim gets absorbed into the + # feature dim (effectively concatenating the camera features). + img_features = einops.rearrange( + img_features, "(b s n) ... -> b s (n ...)", b=batch_size, s=n_obs_steps + ) + global_cond_feats.append(img_features) + + if self.config.env_state_feature: + global_cond_feats.append(batch[OBS_ENV_STATE]) + + # Concatenate features then flatten to (B, global_cond_dim). + return torch.cat(global_cond_feats, dim=-1).flatten(start_dim=1) + + def generate_actions(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: + """ + This function expects `batch` to have: + { + "observation.state": (B, n_obs_steps, state_dim) + + "observation.images": (B, n_obs_steps, num_cameras, C, H, W) + AND/OR + "observation.environment_state": (B, n_obs_steps, environment_dim) + } + """ + batch_size, n_obs_steps = batch[OBS_STATE].shape[:2] + assert n_obs_steps == self.config.n_obs_steps + + # Encode image features and concatenate them all together along with the state vector. + global_cond = self._prepare_global_conditioning(batch) # (B, global_cond_dim) + + # run sampling + actions = self.conditional_sample(batch_size, global_cond=global_cond, noise=noise) + + # Extract `n_action_steps` steps worth of actions (from the current observation). + start = n_obs_steps - 1 + end = start + self.config.n_action_steps + actions = actions[:, start:end] + + return actions + + def compute_loss(self, batch: dict[str, Tensor]) -> Tensor: + """ + This function expects `batch` to have (at least): + { + "observation.state": (B, n_obs_steps, state_dim) + + "observation.images": (B, n_obs_steps, num_cameras, C, H, W) + AND/OR + "observation.environment_state": (B, n_obs_steps, environment_dim) + + "action": (B, horizon, action_dim) + "action_is_pad": (B, horizon) + } + """ + # Input validation. + assert set(batch).issuperset({OBS_STATE, ACTION, "action_is_pad"}) + assert OBS_IMAGES in batch or OBS_ENV_STATE in batch + n_obs_steps = batch[OBS_STATE].shape[1] + horizon = batch[ACTION].shape[1] + assert horizon == self.config.horizon + assert n_obs_steps == self.config.n_obs_steps + + # Encode image features and concatenate them all together along with the state vector. + global_cond = self._prepare_global_conditioning(batch) # (B, global_cond_dim) + + # Forward diffusion. + trajectory = batch[ACTION] + # Sample noise to add to the trajectory. + eps = torch.randn(trajectory.shape, device=trajectory.device) + # Sample a random noising timestep for each item in the batch. + timesteps = torch.randint( + low=0, + high=self.noise_scheduler.config.num_train_timesteps, + size=(trajectory.shape[0],), + device=trajectory.device, + ).long() + # Add noise to the clean trajectories according to the noise magnitude at each timestep. + noisy_trajectory = self.noise_scheduler.add_noise(trajectory, eps, timesteps) + + # Run the denoising network (that might denoise the trajectory, or attempt to predict the noise). + pred = self.unet(noisy_trajectory, timesteps, global_cond=global_cond) + + # Compute the loss. + # The target is either the original trajectory, or the noise. + if self.config.prediction_type == "epsilon": + target = eps + elif self.config.prediction_type == "sample": + target = batch[ACTION] + else: + raise ValueError(f"Unsupported prediction type {self.config.prediction_type}") + + loss = F.mse_loss(pred, target, reduction="none") + + # Mask loss wherever the action is padded with copies (edges of the dataset trajectory). + if self.config.do_mask_loss_for_padding: + if "action_is_pad" not in batch: + raise ValueError( + "You need to provide 'action_is_pad' in the batch when " + f"{self.config.do_mask_loss_for_padding=}." + ) + in_episode_bound = ~batch["action_is_pad"] + loss = loss * in_episode_bound.unsqueeze(-1) + + return loss.mean() + + +class SpatialSoftmax(nn.Module): + """ + Spatial Soft Argmax operation described in "Deep Spatial Autoencoders for Visuomotor Learning" by Finn et al. + (https://huggingface.co/papers/1509.06113). A minimal port of the robomimic implementation. + + At a high level, this takes 2D feature maps (from a convnet/ViT) and returns the "center of mass" + of activations of each channel, i.e., keypoints in the image space for the policy to focus on. + + Example: take feature maps of size (512x10x12). We generate a grid of normalized coordinates (10x12x2): + ----------------------------------------------------- + | (-1., -1.) | (-0.82, -1.) | ... | (1., -1.) | + | (-1., -0.78) | (-0.82, -0.78) | ... | (1., -0.78) | + | ... | ... | ... | ... | + | (-1., 1.) | (-0.82, 1.) | ... | (1., 1.) | + ----------------------------------------------------- + This is achieved by applying channel-wise softmax over the activations (512x120) and computing the dot + product with the coordinates (120x2) to get expected points of maximal activation (512x2). + + The example above results in 512 keypoints (corresponding to the 512 input channels). We can optionally + provide num_kp != None to control the number of keypoints. This is achieved by a first applying a learnable + linear mapping (in_channels, H, W) -> (num_kp, H, W). + """ + + def __init__(self, input_shape, num_kp=None): + """ + Args: + input_shape (list): (C, H, W) input feature map shape. + num_kp (int): number of keypoints in output. If None, output will have the same number of channels as input. + """ + super().__init__() + + assert len(input_shape) == 3 + self._in_c, self._in_h, self._in_w = input_shape + + if num_kp is not None: + self.nets = torch.nn.Conv2d(self._in_c, num_kp, kernel_size=1) + self._out_c = num_kp + else: + self.nets = None + self._out_c = self._in_c + + # we could use torch.linspace directly but that seems to behave slightly differently than numpy + # and causes a small degradation in pc_success of pre-trained models. + pos_x, pos_y = np.meshgrid(np.linspace(-1.0, 1.0, self._in_w), np.linspace(-1.0, 1.0, self._in_h)) + pos_x = torch.from_numpy(pos_x.reshape(self._in_h * self._in_w, 1)).float() + pos_y = torch.from_numpy(pos_y.reshape(self._in_h * self._in_w, 1)).float() + # register as buffer so it's moved to the correct device. + self.register_buffer("pos_grid", torch.cat([pos_x, pos_y], dim=1)) + + def forward(self, features: Tensor) -> Tensor: + """ + Args: + features: (B, C, H, W) input feature maps. + Returns: + (B, K, 2) image-space coordinates of keypoints. + """ + if self.nets is not None: + features = self.nets(features) + + # [B, K, H, W] -> [B * K, H * W] where K is number of keypoints + features = features.reshape(-1, self._in_h * self._in_w) + # 2d softmax normalization + attention = F.softmax(features, dim=-1) + # [B * K, H * W] x [H * W, 2] -> [B * K, 2] for spatial coordinate mean in x and y dimensions + expected_xy = attention @ self.pos_grid + # reshape to [B, K, 2] + feature_keypoints = expected_xy.view(-1, self._out_c, 2) + + return feature_keypoints + + +class DiffusionRgbEncoder(nn.Module): + """Encodes an RGB image into a 1D feature vector. + + Includes the ability to normalize and crop the image first. + """ + + def __init__(self, config: DiffusionConfig): + super().__init__() + # Set up optional preprocessing. + if config.crop_shape is not None: + self.do_crop = True + # Always use center crop for eval + self.center_crop = torchvision.transforms.CenterCrop(config.crop_shape) + if config.crop_is_random: + self.maybe_random_crop = torchvision.transforms.RandomCrop(config.crop_shape) + else: + self.maybe_random_crop = self.center_crop + else: + self.do_crop = False + + # Set up backbone. + backbone_model = getattr(torchvision.models, config.vision_backbone)( + weights=config.pretrained_backbone_weights + ) + # Note: This assumes that the layer4 feature map is children()[-3] + # TODO(alexander-soare): Use a safer alternative. + self.backbone = nn.Sequential(*(list(backbone_model.children())[:-2])) + if config.use_group_norm: + if config.pretrained_backbone_weights: + raise ValueError( + "You can't replace BatchNorm in a pretrained model without ruining the weights!" + ) + self.backbone = _replace_submodules( + root_module=self.backbone, + predicate=lambda x: isinstance(x, nn.BatchNorm2d), + func=lambda x: nn.GroupNorm(num_groups=x.num_features // 16, num_channels=x.num_features), + ) + + # Set up pooling and final layers. + # Use a dry run to get the feature map shape. + # The dummy input should take the number of image channels from `config.image_features` and it should + # use the height and width from `config.crop_shape` if it is provided, otherwise it should use the + # height and width from `config.image_features`. + + # Note: we have a check in the config class to make sure all images have the same shape. + images_shape = next(iter(config.image_features.values())).shape + dummy_shape_h_w = config.crop_shape if config.crop_shape is not None else images_shape[1:] + dummy_shape = (1, images_shape[0], *dummy_shape_h_w) + feature_map_shape = get_output_shape(self.backbone, dummy_shape)[1:] + + self.pool = SpatialSoftmax(feature_map_shape, num_kp=config.spatial_softmax_num_keypoints) + self.feature_dim = config.spatial_softmax_num_keypoints * 2 + self.out = nn.Linear(config.spatial_softmax_num_keypoints * 2, self.feature_dim) + self.relu = nn.ReLU() + + def forward(self, x: Tensor) -> Tensor: + """ + Args: + x: (B, C, H, W) image tensor with pixel values in [0, 1]. + Returns: + (B, D) image feature. + """ + # Preprocess: maybe crop (if it was set up in the __init__). + if self.do_crop: + if self.training: # noqa: SIM108 + x = self.maybe_random_crop(x) + else: + # Always use center crop for eval. + x = self.center_crop(x) + # Extract backbone feature. + x = torch.flatten(self.pool(self.backbone(x)), start_dim=1) + # Final linear layer with non-linearity. + x = self.relu(self.out(x)) + return x + + +def _replace_submodules( + root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module] +) -> nn.Module: + """ + Args: + root_module: The module for which the submodules need to be replaced + predicate: Takes a module as an argument and must return True if the that module is to be replaced. + func: Takes a module as an argument and returns a new module to replace it with. + Returns: + The root module with its submodules replaced. + """ + if predicate(root_module): + return func(root_module) + + replace_list = [k.split(".") for k, m in root_module.named_modules(remove_duplicate=True) if predicate(m)] + for *parents, k in replace_list: + parent_module = root_module + if len(parents) > 0: + parent_module = root_module.get_submodule(".".join(parents)) + if isinstance(parent_module, nn.Sequential): + src_module = parent_module[int(k)] + else: + src_module = getattr(parent_module, k) + tgt_module = func(src_module) + if isinstance(parent_module, nn.Sequential): + parent_module[int(k)] = tgt_module + else: + setattr(parent_module, k, tgt_module) + # verify that all BN are replaced + assert not any(predicate(m) for _, m in root_module.named_modules(remove_duplicate=True)) + return root_module + + +class DiffusionSinusoidalPosEmb(nn.Module): + """1D sinusoidal positional embeddings as in Attention is All You Need.""" + + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, x: Tensor) -> Tensor: + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x.unsqueeze(-1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class DiffusionConv1dBlock(nn.Module): + """Conv1d --> GroupNorm --> Mish""" + + def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): + super().__init__() + + self.block = nn.Sequential( + nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2), + nn.GroupNorm(n_groups, out_channels), + nn.Mish(), + ) + + def forward(self, x): + return self.block(x) + + +class DiffusionConditionalUnet1d(nn.Module): + """A 1D convolutional UNet with FiLM modulation for conditioning. + + Note: this removes local conditioning as compared to the original diffusion policy code. + """ + + def __init__(self, config: DiffusionConfig, global_cond_dim: int): + super().__init__() + + self.config = config + + # Encoder for the diffusion timestep. + self.diffusion_step_encoder = nn.Sequential( + DiffusionSinusoidalPosEmb(config.diffusion_step_embed_dim), + nn.Linear(config.diffusion_step_embed_dim, config.diffusion_step_embed_dim * 4), + nn.Mish(), + nn.Linear(config.diffusion_step_embed_dim * 4, config.diffusion_step_embed_dim), + ) + + # The FiLM conditioning dimension. + cond_dim = config.diffusion_step_embed_dim + global_cond_dim + + # In channels / out channels for each downsampling block in the Unet's encoder. For the decoder, we + # just reverse these. + in_out = [(config.action_feature.shape[0], config.down_dims[0])] + list( + zip(config.down_dims[:-1], config.down_dims[1:], strict=True) + ) + + # Unet encoder. + common_res_block_kwargs = { + "cond_dim": cond_dim, + "kernel_size": config.kernel_size, + "n_groups": config.n_groups, + "use_film_scale_modulation": config.use_film_scale_modulation, + } + self.down_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (len(in_out) - 1) + self.down_modules.append( + nn.ModuleList( + [ + DiffusionConditionalResidualBlock1d(dim_in, dim_out, **common_res_block_kwargs), + DiffusionConditionalResidualBlock1d(dim_out, dim_out, **common_res_block_kwargs), + # Downsample as long as it is not the last block. + nn.Conv1d(dim_out, dim_out, 3, 2, 1) if not is_last else nn.Identity(), + ] + ) + ) + + # Processing in the middle of the auto-encoder. + self.mid_modules = nn.ModuleList( + [ + DiffusionConditionalResidualBlock1d( + config.down_dims[-1], config.down_dims[-1], **common_res_block_kwargs + ), + DiffusionConditionalResidualBlock1d( + config.down_dims[-1], config.down_dims[-1], **common_res_block_kwargs + ), + ] + ) + + # Unet decoder. + self.up_modules = nn.ModuleList([]) + for ind, (dim_out, dim_in) in enumerate(reversed(in_out[1:])): + is_last = ind >= (len(in_out) - 1) + self.up_modules.append( + nn.ModuleList( + [ + # dim_in * 2, because it takes the encoder's skip connection as well + DiffusionConditionalResidualBlock1d(dim_in * 2, dim_out, **common_res_block_kwargs), + DiffusionConditionalResidualBlock1d(dim_out, dim_out, **common_res_block_kwargs), + # Upsample as long as it is not the last block. + nn.ConvTranspose1d(dim_out, dim_out, 4, 2, 1) if not is_last else nn.Identity(), + ] + ) + ) + + self.final_conv = nn.Sequential( + DiffusionConv1dBlock(config.down_dims[0], config.down_dims[0], kernel_size=config.kernel_size), + nn.Conv1d(config.down_dims[0], config.action_feature.shape[0], 1), + ) + + def forward(self, x: Tensor, timestep: Tensor | int, global_cond=None) -> Tensor: + """ + Args: + x: (B, T, input_dim) tensor for input to the Unet. + timestep: (B,) tensor of (timestep_we_are_denoising_from - 1). + global_cond: (B, global_cond_dim) + output: (B, T, input_dim) + Returns: + (B, T, input_dim) diffusion model prediction. + """ + # For 1D convolutions we'll need feature dimension first. + x = einops.rearrange(x, "b t d -> b d t") + + timesteps_embed = self.diffusion_step_encoder(timestep) + + # If there is a global conditioning feature, concatenate it to the timestep embedding. + if global_cond is not None: + global_feature = torch.cat([timesteps_embed, global_cond], axis=-1) + else: + global_feature = timesteps_embed + + # Run encoder, keeping track of skip features to pass to the decoder. + encoder_skip_features: list[Tensor] = [] + for resnet, resnet2, downsample in self.down_modules: + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + encoder_skip_features.append(x) + x = downsample(x) + + for mid_module in self.mid_modules: + x = mid_module(x, global_feature) + + # Run decoder, using the skip features from the encoder. + for resnet, resnet2, upsample in self.up_modules: + x = torch.cat((x, encoder_skip_features.pop()), dim=1) + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + x = upsample(x) + + x = self.final_conv(x) + + x = einops.rearrange(x, "b d t -> b t d") + return x + + +class DiffusionConditionalResidualBlock1d(nn.Module): + """ResNet style 1D convolutional block with FiLM modulation for conditioning.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + cond_dim: int, + kernel_size: int = 3, + n_groups: int = 8, + # Set to True to do scale modulation with FiLM as well as bias modulation (defaults to False meaning + # FiLM just modulates bias). + use_film_scale_modulation: bool = False, + ): + super().__init__() + + self.use_film_scale_modulation = use_film_scale_modulation + self.out_channels = out_channels + + self.conv1 = DiffusionConv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups) + + # FiLM modulation (https://huggingface.co/papers/1709.07871) outputs per-channel bias and (maybe) scale. + cond_channels = out_channels * 2 if use_film_scale_modulation else out_channels + self.cond_encoder = nn.Sequential(nn.Mish(), nn.Linear(cond_dim, cond_channels)) + + self.conv2 = DiffusionConv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups) + + # A final convolution for dimension matching the residual (if needed). + self.residual_conv = ( + nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity() + ) + + def forward(self, x: Tensor, cond: Tensor) -> Tensor: + """ + Args: + x: (B, in_channels, T) + cond: (B, cond_dim) + Returns: + (B, out_channels, T) + """ + out = self.conv1(x) + + # Get condition embedding. Unsqueeze for broadcasting to `out`, resulting in (B, out_channels, 1). + cond_embed = self.cond_encoder(cond).unsqueeze(-1) + if self.use_film_scale_modulation: + # Treat the embedding as a list of scales and biases. + scale = cond_embed[:, : self.out_channels] + bias = cond_embed[:, self.out_channels :] + out = scale * out + bias + else: + # Treat the embedding as biases. + out = out + cond_embed + + out = self.conv2(out) + out = out + self.residual_conv(x) + return out diff --git a/diffusion/processor_diffusion.py b/diffusion/processor_diffusion.py new file mode 100644 index 0000000..a7799be --- /dev/null +++ b/diffusion/processor_diffusion.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python + +# Copyright 2024 Columbia Artificial Intelligence, Robotics Lab, +# and The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Any + +import torch + +from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig +from lerobot.processor import ( + AddBatchDimensionProcessorStep, + DeviceProcessorStep, + NormalizerProcessorStep, + PolicyAction, + PolicyProcessorPipeline, + RenameObservationsProcessorStep, + UnnormalizerProcessorStep, +) +from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action +from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME + + +def make_diffusion_pre_post_processors( + config: DiffusionConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """ + Constructs pre-processor and post-processor pipelines for a diffusion policy. + + The pre-processing pipeline prepares the input data for the model by: + 1. Renaming features. + 2. Normalizing the input and output features based on dataset statistics. + 3. Adding a batch dimension. + 4. Moving the data to the specified device. + + The post-processing pipeline handles the model's output by: + 1. Moving the data to the CPU. + 2. Unnormalizing the output features to their original scale. + + Args: + config: The configuration object for the diffusion policy, + containing feature definitions, normalization mappings, and device information. + dataset_stats: A dictionary of statistics used for normalization. + Defaults to None. + + Returns: + A tuple containing the configured pre-processor and post-processor pipelines. + """ + + input_steps = [ + RenameObservationsProcessorStep(rename_map={}), + AddBatchDimensionProcessorStep(), + DeviceProcessorStep(device=config.device), + NormalizerProcessorStep( + features={**config.input_features, **config.output_features}, + norm_map=config.normalization_mapping, + stats=dataset_stats, + ), + ] + output_steps = [ + UnnormalizerProcessorStep( + features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats + ), + DeviceProcessorStep(device="cpu"), + ] + return ( + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( + steps=input_steps, + name=POLICY_PREPROCESSOR_DEFAULT_NAME, + ), + PolicyProcessorPipeline[PolicyAction, PolicyAction]( + steps=output_steps, + name=POLICY_POSTPROCESSOR_DEFAULT_NAME, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ), + ) diff --git a/docs/superpowers/plans/2026-03-30-streaming-hdf5-ee-action.md b/docs/superpowers/plans/2026-03-30-streaming-hdf5-ee-action.md new file mode 100644 index 0000000..1e697c1 --- /dev/null +++ b/docs/superpowers/plans/2026-03-30-streaming-hdf5-ee-action.md @@ -0,0 +1,42 @@ +# Streaming HDF5 EE Action Dataset Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 将 Diana 仿真采集改为流式写入 HDF5，图像保存为 256x256 的四路相机视角，并把 `/action` 改为 IK 前的原始末端位姿动作。 + +**Architecture:** 新增一个独立的流式 HDF5 episode writer，负责逐帧写入 qpos、原始 action 和 resize 后图像，并在 episode 成功时原子提交、失败时删除临时文件。采集脚本只负责 rollout 和把每一步观测/动作交给 writer，避免整集数据先堆在内存里。 + +**Tech Stack:** Python, h5py, numpy, cv2, unittest, MuJoCo demo scripts + +--- + +### Task 1: 为流式 writer 建立测试边界 + +**Files:** +- Create: `tests/test_streaming_episode_writer.py` +- Create: `roboimi/utils/streaming_episode_writer.py` + +- [ ] **Step 1: Write the failing test** +- [ ] **Step 2: Run `python -m unittest tests.test_streaming_episode_writer -v` and confirm it fails because the writer module does not exist** +- [ ] **Step 3: Implement the minimal streaming writer with temp-file commit/discard, per-frame append, and 256x256 image resize** +- [ ] **Step 4: Re-run `python -m unittest tests.test_streaming_episode_writer -v` and confirm it passes** + +### Task 2: 接入 Diana 采集脚本 + +**Files:** +- Modify: `roboimi/demos/diana_record_sim_episodes.py` +- Reuse: `roboimi/utils/streaming_episode_writer.py` + +- [ ] **Step 1: Replace in-memory `data_dict` / `obs` accumulation with per-episode streaming writer lifecycle** +- [ ] **Step 2: Keep four cameras (`angle`, `r_vis`, `top`, `front`) and resize to 256x256 before persistence** +- [ ] **Step 3: Capture raw policy output before IK and write that to `/action`** +- [ ] **Step 4: On success commit to `episode_{idx}.hdf5`; on failure remove temp file** + +### Task 3: 验证改动 + +**Files:** +- Verify only + +- [ ] **Step 1: Run unit tests for the writer** +- [ ] **Step 2: Run one end-to-end collection episode and stop after `episode_0.hdf5` becomes readable** +- [ ] **Step 3: Verify HDF5 keys and shapes: `action=(700,16)`, image datasets are `(700,256,256,3)`, and `/action` matches raw EE action semantics** diff --git a/docs/superpowers/plans/2026-03-31-raw-action-trajectory-viewer.md b/docs/superpowers/plans/2026-03-31-raw-action-trajectory-viewer.md new file mode 100644 index 0000000..93b42c7 --- /dev/null +++ b/docs/superpowers/plans/2026-03-31-raw-action-trajectory-viewer.md @@ -0,0 +1,26 @@ +# Raw Action Trajectory Viewer Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 在可交互 MuJoCo 仿真窗口中，把 rollout 导出的 raw EE action 轨迹用红色轨迹标出来并启动仿真供人工查看。 + +**Architecture:** 读取已有 trajectory artifact 中的 raw_action / step 数据，生成左右臂末端轨迹点，并在 viewer 渲染循环中持续注入红色 marker。实现尽量独立为一个可复用的小脚本，避免影响训练/评估主路径。 + +**Tech Stack:** Python, NumPy, MuJoCo viewer, unittest/mock. + +--- + +### Task 1: 抽取 raw_action 轨迹并生成可视化点集 +- [ ] 写失败测试，验证从 trajectory.npz 提取左右臂轨迹点 +- [ ] 实现最小 helper +- [ ] 运行测试确认通过 + +### Task 2: 在 viewer 中渲染红色轨迹并支持交互查看 +- [ ] 写失败测试，验证 marker 配置/调用 +- [ ] 实现 viewer 可视化脚本 +- [ ] 运行测试确认通过 + +### Task 3: 启动真实仿真窗口供人工查看 +- [ ] 用现有 trajectory artifact 启动 viewer +- [ ] 确认窗口可交互、红线出现 +- [ ] 向用户汇报启动方式与脚本路径 diff --git a/docs/superpowers/plans/2026-03-31-rollout-artifacts.md b/docs/superpowers/plans/2026-03-31-rollout-artifacts.md new file mode 100644 index 0000000..00f5aa8 --- /dev/null +++ b/docs/superpowers/plans/2026-03-31-rollout-artifacts.md @@ -0,0 +1,44 @@ +# Rollout Artifacts Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Extend rollout evaluation so one selected checkpoint can be run once with video capture, timing breakdown, and saved EE trajectory artifacts. + +**Architecture:** Keep the implementation centered in `eval_vla.py` so existing training-time rollout validation remains compatible. Add config-gated artifact capture helpers, serialize outputs under the eval run directory, and add lightweight tests for helper behavior and summary wiring; default eval behavior must remain unchanged when artifact capture is off. + +**Tech Stack:** Python, Hydra/OmegaConf, NumPy, OpenCV, JSON, PyTorch unittest/mocking. + +--- + +### Task 1: Add artifact capture configuration and helper wiring + +**Files:** +- Modify: `roboimi/demos/vla_scripts/eval_vla.py` +- Modify: `roboimi/vla/conf/eval/eval.yaml` +- Test: `tests/test_eval_vla_rollout_artifacts.py` + +- [ ] **Step 1: Write failing tests for optional artifact config / summary wiring** +- [ ] **Step 2: Implement config-backed artifact flags and output paths with defaults that write nothing** +- [ ] **Step 3: Verify existing eval call sites still work with defaults** + +### Task 2: Add timing breakdown, video recording, and trajectory export + +**Files:** +- Modify: `roboimi/demos/vla_scripts/eval_vla.py` +- Test: `tests/test_eval_vla_rollout_artifacts.py` + +- [ ] **Step 1: Write failing tests for timing aggregation, trajectory serialization, and summary schema** +- [ ] **Step 2: Implement per-step timing capture for `obs_read_ms`, `preprocess_ms`, `inference_ms`, `env_step_ms`, `loop_total_ms`** +- [ ] **Step 3: Implement MP4 recording from a chosen camera stream and canonical `trajectory.npz` export using `left_link7/right_link7` executed poses after `env.step`** +- [ ] **Step 4: Run focused tests and fix issues** + +### Task 3: Stop training safely and execute one real rollout + +**Files:** +- Use: `roboimi/demos/vla_scripts/eval_vla.py` +- Output: `runs/.../eval_artifacts/...` + +- [ ] **Step 1: Stop the active training process, wait for exit, and confirm the target checkpoint is readable** +- [ ] **Step 2: Select the latest completed checkpoint if an explicit one is not provided; fall back to prior completed / best checkpoint if needed** +- [ ] **Step 3: Run one headless rollout with artifact capture enabled** +- [ ] **Step 4: Verify the MP4 / timing summary / trajectory files exist and summarize findings** diff --git a/docs/superpowers/plans/2026-04-01-imf-attnres-policy-migration.md b/docs/superpowers/plans/2026-04-01-imf-attnres-policy-migration.md new file mode 100644 index 0000000..97a8d7c --- /dev/null +++ b/docs/superpowers/plans/2026-04-01-imf-attnres-policy-migration.md @@ -0,0 +1,268 @@ +# IMF-AttnRes Policy Migration Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 将 external `diffusion_policy@185ed659` 的 IMF-AttnRes 模型、训练目标和一步推理机制迁移到 RoboIMI，并在保持三相机视觉条件输入与现有训练/rollout 工作流的前提下启动同参数训练。 + +**Architecture:** 保留 RoboIMI 现有 ResNet 三相机观测编码、normalization、queue-based online rollout 和训练脚本；新增 AttnRes 组件与 IMF transformer head，并新增 IMF 专用 agent 以覆盖 DDPM loss / DDIM inference 语义。训练脚本只做最小接线修改，让新 head/agent 能用现有 optimizer、checkpoint、SwanLab 和 headless rollout。 + +**Tech Stack:** PyTorch, Hydra, diffusers schedulers (仅保留兼容初始化), MuJoCo rollout, unittest, SwanLab + +--- + +## File Map + +### New files +- `roboimi/vla/models/heads/attnres_transformer_components.py` — 本地 IMF AttnRes 基础组件 +- `roboimi/vla/models/heads/imf_transformer1d.py` — IMF transformer head，暴露 `forward(sample, r, t, cond=None)` +- `roboimi/vla/agent_imf.py` — IMF 专用 VLA agent，复用现有观测/队列/normalization 逻辑并覆盖 loss / inference +- `roboimi/vla/conf/head/imf_transformer1d.yaml` — IMF head 配置 +- `roboimi/vla/conf/agent/resnet_imf_attnres.yaml` — IMF agent + backbone/head 组合配置 +- `tests/test_imf_transformer1d_external_alignment.py` — external `185ed659` 对齐测试 +- `tests/test_imf_vla_agent.py` — IMF agent 的 loss / inference / queue 语义测试 + +### Modified files +- `roboimi/demos/vla_scripts/train_vla.py` — 优化器参数分组接线；确保新 agent 能无缝训练 +- `roboimi/vla/conf/config.yaml` — 保持默认配置不变，仅支持通过 override 启用 IMF agent +- `tests/test_train_vla_transformer_optimizer.py` — 覆盖 IMF head 的 optimizer-group 行为 +- （如需要）`roboimi/vla/models/heads/__init__.py` 或相近导出文件 — 暴露新 head + +--- + +### Task 1: 写 IMF transformer 对齐测试 + +**Files:** +- Create: `tests/test_imf_transformer1d_external_alignment.py` +- Reference: `/home/droid/project/diffusion_policy/diffusion_policy/model/diffusion/attnres_transformer_components.py` +- Reference: `/home/droid/project/diffusion_policy/diffusion_policy/model/diffusion/imf_transformer_for_diffusion.py` + +- [ ] **Step 1: 写失败测试，验证 local IMF head 与 external `185ed659` 的 state-dict key、前向 shape、forward 数值、optim groups 对齐** + +```python +with torch.no_grad(): + external_out = external_model(sample=sample, r=r, t=t, cond=cond) + local_out = local_model(sample=sample, r=r, t=t, cond=cond) +assert torch.allclose(local_out, external_out, atol=1e-6, rtol=1e-5) +``` + +- [ ] **Step 2: 运行单测，确认当前失败** + +Run: `python -m unittest tests.test_imf_transformer1d_external_alignment -v` +Expected: FAIL，提示 `imf_transformer1d` / `attnres` 模块不存在 + +- [ ] **Step 3: 若测试需要复用现有 external-loader 逻辑，则从 `tests/test_transformer1d_external_alignment.py` 复制最小必要 helper，避免重复依赖 session context** + +- [ ] **Step 4: 提交测试骨架** + +```bash +git add tests/test_imf_transformer1d_external_alignment.py +git commit -m "test: add IMF transformer external alignment coverage" +``` + +### Task 2: 实现 AttnRes 组件与 IMF transformer head + +**Files:** +- Create: `roboimi/vla/models/heads/attnres_transformer_components.py` +- Create: `roboimi/vla/models/heads/imf_transformer1d.py` +- Modify: `tests/test_imf_transformer1d_external_alignment.py` + +- [ ] **Step 1: 按 external `185ed659` 迁移 AttnRes 基础组件，保持命名和参数语义一致** + +必须包含： +- `RMSNorm` +- `RMSNormNoWeight` +- `precompute_rope_freqs` +- `apply_rope` +- `GroupedQuerySelfAttention` +- `SwiGLUFFN` +- `AttnResOperator` +- `AttnResSubLayer` +- `AttnResTransformerBackbone` + +- [ ] **Step 2: 在 `imf_transformer1d.py` 中实现本地 IMF head** + +必须满足： +- `forward(sample, r, t, cond=None)` +- 默认支持 `backbone_type='attnres_full'` +- token 序列为 `[r_token, t_token, cond_tokens..., sample_tokens...]` +- 输出只切回 sample token 段 +- 保留 `get_optim_groups()` 供 AdamW 分组 + +- [ ] **Step 3: 运行对齐测试，修正 state-dict key / init / no-decay 参数分组不一致问题** + +Run: `python -m unittest tests.test_imf_transformer1d_external_alignment -v` +Expected: PASS + +- [ ] **Step 4: 提交模型组件实现** + +```bash +git add roboimi/vla/models/heads/attnres_transformer_components.py \ + roboimi/vla/models/heads/imf_transformer1d.py \ + tests/test_imf_transformer1d_external_alignment.py +git commit -m "feat: add IMF AttnRes transformer head" +``` + +### Task 3: 写 IMF agent 行为测试 + +**Files:** +- Create: `tests/test_imf_vla_agent.py` +- Reference: `roboimi/vla/agent.py` +- Reference: `tests/test_resnet_transformer_agent_wiring.py` + +- [ ] **Step 1: 写失败测试，覆盖 IMF agent 的核心契约** + +需要覆盖： +1. `compute_loss()` 接受当前 batch 结构并返回标量 loss +2. `predict_action()` 输出 `(B, pred_horizon, action_dim)` +3. `select_action()` 仍按 queue/chunk 语义工作 +4. `predict_action()` 不走 DDIM 多步循环，而是只触发一步 IMF sample +5. `action_is_pad` 存在时仅在有效 action 上计 loss + +- [ ] **Step 2: 用 stub backbone / stub head 记录调用参数，验证 `r,t,cond` 的传递与 observation conditioning 维度正确** + +```python +self.assertEqual(recorded['cond'].shape, (B, obs_horizon, expected_cond_dim)) +self.assertTrue(torch.allclose(recorded['r'], torch.zeros(B))) +self.assertTrue(torch.allclose(recorded['t'], torch.ones(B))) +``` + +- [ ] **Step 3: 运行测试，确认当前失败** + +Run: `python -m unittest tests.test_imf_vla_agent -v` +Expected: FAIL，提示 `roboimi.vla.agent_imf` 不存在 + +- [ ] **Step 4: 提交测试骨架** + +```bash +git add tests/test_imf_vla_agent.py +git commit -m "test: add IMF VLA agent behavior coverage" +``` + +### Task 4: 实现 IMF agent 与 Hydra 接线 + +**Files:** +- Create: `roboimi/vla/agent_imf.py` +- Create: `roboimi/vla/conf/head/imf_transformer1d.yaml` +- Create: `roboimi/vla/conf/agent/resnet_imf_attnres.yaml` +- Modify: `roboimi/demos/vla_scripts/train_vla.py` +- Modify: `tests/test_train_vla_transformer_optimizer.py` +- Modify: `tests/test_imf_vla_agent.py` + +- [ ] **Step 1: 以 `VLAAgent` 为基础实现 `IMFVLAAgent`** + +实现策略： +- 复用 `VLAAgent.__init__`、`_build_cond()`、`reset()`、`_populate_queues()`、`_prepare_observation_batch()`、`select_action()`、`get_normalization_stats()` +- 覆盖： + - `compute_loss()` -> IMF objective + - `predict_action()` -> one-step sample +- 提供内部 helper： + - `_broadcast_batch_time` + - `_apply_conditioning`（如需） + - `_compute_u_and_du_dt` + - `_compound_velocity` + - `_sample_one_step` + +- [ ] **Step 2: 在 JVP 路径中加入 CUDA math SDPA fallback，保持 external repo 的稳定性策略** + +- [ ] **Step 3: 新增 Hydra 配置，让 `agent=resnet_imf_attnres` 可实例化** + +关键默认值： +- `_target_: roboimi.vla.agent_imf.IMFVLAAgent` +- `head._target_: roboimi.vla.models.heads.imf_transformer1d.IMFTransformer1D` +- `head.backbone_type: attnres_full` +- `head.causal_attn: false` +- `head.time_as_cond: true` +- `head.n_cond_layers: 0` +- `inference_steps: 1` +- `camera_names: ${data.camera_names}` +- `vision_backbone.camera_names: ${agent.camera_names}` + +- [ ] **Step 4: 让训练脚本对任何带 `get_optim_groups()` 的 head 复用参数分组，而不是硬编码旧 transformer head_type** + +推荐最小改法： +```python +use_head_groups = callable(getattr(noise_pred_net, 'get_optim_groups', None)) +``` + +- [ ] **Step 5: 运行测试并修复 wiring 问题** + +Run: +- `python -m unittest tests.test_imf_vla_agent -v` +- `python -m unittest tests.test_train_vla_transformer_optimizer -v` + +Expected: PASS + +- [ ] **Step 6: 提交 agent / config / train-script 接线** + +```bash +git add roboimi/vla/agent_imf.py \ + roboimi/vla/conf/head/imf_transformer1d.yaml \ + roboimi/vla/conf/agent/resnet_imf_attnres.yaml \ + roboimi/demos/vla_scripts/train_vla.py \ + tests/test_imf_vla_agent.py \ + tests/test_train_vla_transformer_optimizer.py +git commit -m "feat: add IMF VLA agent and training wiring" +``` + +### Task 5: 集成验证与训练启动 + +**Files:** +- Modify: none required unless验证暴露真实问题 +- Use run artifacts under: `runs/` + +- [ ] **Step 1: 运行聚焦测试集** + +Run: +```bash +python -m unittest \ + tests.test_imf_transformer1d_external_alignment \ + tests.test_imf_vla_agent \ + tests.test_resnet_transformer_agent_wiring \ + tests.test_train_vla_transformer_optimizer -v +``` +Expected: PASS + +- [ ] **Step 2: 运行一个最小 GPU 训练冒烟任务（不必长跑）** + +Run: +```bash +/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \ + agent=resnet_imf_attnres \ + data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \ + data.camera_names=[r_vis,top,front] \ + train.device=cuda train.max_steps=2 train.batch_size=4 train.num_workers=2 \ + train.use_swanlab=false train.rollout_val_freq_epochs=0 +``` +Expected: 成功完成 2 steps，生成 checkpoint / log，无 shape 或 JVP 错误 + +- [ ] **Step 3: 用正式参数启动 IMF 训练** + +Run: +```bash +/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \ + agent=resnet_imf_attnres \ + data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \ + data.camera_names=[r_vis,top,front] \ + train.device=cuda train.val_split=0.0 train.seed=42 \ + train.batch_size=80 train.lr=5e-4 train.num_workers=12 train.max_steps=150000 \ + train.log_freq=100 train.save_freq=10000 train.use_swanlab=true \ + train.swanlab_project=roboimi-vla \ + train.rollout_val_freq_epochs=5 train.rollout_validate_on_checkpoint=false \ + train.rollout_num_episodes=5 train.warmup_steps=2000 \ + train.scheduler_type=cosine train.min_lr=1e-6 train.weight_decay=1e-5 train.grad_clip=1.0 \ + agent.pred_horizon=16 agent.inference_steps=1 \ + agent.head.n_emb=384 agent.head.n_layer=18 agent.head.n_head=1 agent.head.n_kv_head=1 \ + agent.vision_backbone.pretrained_backbone_weights=null \ + agent.vision_backbone.freeze_backbone=false \ + agent.vision_backbone.use_separate_rgb_encoder_per_camera=true +``` +Expected: 训练启动成功，SwanLab 记录完整 config，5 epoch 一次 headless rollout + +- [ ] **Step 4: 记录 run 路径、训练 PID、SwanLab 运行名并向用户汇报** + +- [ ] **Step 5: 提交最终收尾改动（如果 smoke fix 需要额外 patch）** + +```bash +git add +git commit -m "chore: verify IMF AttnRes training launch" +``` diff --git a/docs/superpowers/plans/2026-04-02-imf-rollout-trajectory-images-and-short-horizon-training.md b/docs/superpowers/plans/2026-04-02-imf-rollout-trajectory-images-and-short-horizon-training.md new file mode 100644 index 0000000..9611ab9 --- /dev/null +++ b/docs/superpowers/plans/2026-04-02-imf-rollout-trajectory-images-and-short-horizon-training.md @@ -0,0 +1,79 @@ +# IMF Rollout Trajectory Images and Short-Horizon Training Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add training-time rollout front trajectory image export plus SwanLab image logging, then start a new local IMF training run with `emb=384`, `layer=12`, `pred_horizon=8`, `num_action_steps=4`, `max_steps=50000`. + +**Architecture:** Extend `eval_vla.py` so a rollout can emit one per-episode static front-view image with red EE trajectory overlay. Extend `train_vla.py` so rollout validation forces image export, forces video off, and uploads those per-episode images to SwanLab. Launch the requested new run through explicit command-line overrides rather than branch-default config changes. + +**Tech Stack:** Python, PyTorch, Hydra/OmegaConf, MuJoCo, OpenCV, SwanLab. + +--- + +### Task 1: Add and validate rollout image tests + +**Files:** +- Modify: `tests/test_eval_vla_rollout_artifacts.py` +- Modify: `tests/test_train_vla_swanlab_logging.py` +- Modify: `tests/test_train_vla_rollout_validation.py` + +- [ ] Add/adjust eval tests so they assert per-episode trajectory image paths are produced without requiring video export. +- [ ] Add/adjust training tests so they assert training-time rollout validation forces `record_video=false`. +- [ ] Add/adjust training tests so they assert trajectory image paths flow from eval summary into SwanLab media logging. +- [ ] Add/adjust training tests so they assert image media is logged, not only scalar reward metrics. + +### Task 2: Implement per-episode front trajectory image export in eval + +**Files:** +- Modify: `roboimi/demos/vla_scripts/eval_vla.py` +- Reuse/Read: `roboimi/utils/raw_action_trajectory_viewer.py` +- Modify: `roboimi/vla/conf/eval/eval.yaml` + +- [ ] Add config plumbing for `save_trajectory_image` and `trajectory_image_camera_name`. +- [ ] Ensure the default training-time camera resolution path is pinned to `front`. +- [ ] Implement distinct per-episode image naming so 5 rollout episodes create 5 distinct PNGs. +- [ ] Reuse the existing red trajectory representation logic when composing the PNG. +- [ ] Ensure headless eval works under EGL even on machines with `DISPLAY` set. + +### Task 3: Implement SwanLab rollout image logging in training + +**Files:** +- Modify: `roboimi/demos/vla_scripts/train_vla.py` +- Modify: `tests/test_train_vla_swanlab_logging.py` +- Modify: `tests/test_train_vla_rollout_validation.py` + +- [ ] Make `run_rollout_validation()` force `record_video=false`. +- [ ] Make `run_rollout_validation()` force `save_trajectory_image=true` and `trajectory_image_camera_name=front`. +- [ ] Ensure rollout validation still uses 5 episodes per validation event for the requested run. +- [ ] Add a best-effort helper that converts per-episode image paths into SwanLab image media payloads. +- [ ] Keep image-upload failures non-fatal and warning-only. + +### Task 4: Verify action-chunk semantics for the new run + +**Files:** +- Verify: `roboimi/vla/agent.py` +- Verify: `roboimi/vla/agent_imf.py` +- Test: `tests/test_imf_vla_agent.py` + +- [ ] Confirm the existing queue logic still means “predict 8, execute first 4”. +- [ ] Do not change branch defaults unless strictly necessary; prefer launch-time overrides. + +### Task 5: Verify and launch the requested local training run + +**Files:** +- Use: `roboimi/demos/vla_scripts/train_vla.py` +- Use: `roboimi/demos/vla_scripts/eval_vla.py` + +- [ ] Run the targeted verification suite. +- [ ] Run one real headless smoke eval and confirm a front trajectory PNG is produced while `video_mp4` stays null. +- [ ] Launch the new local training run with explicit overrides including: + - `agent=resnet_imf_attnres` + - `agent.head.n_emb=384` + - `agent.head.n_layer=12` + - `agent.pred_horizon=8` + - `agent.num_action_steps=4` + - `train.max_steps=50000` + - `train.rollout_num_episodes=5` + - `train.use_swanlab=true` + - current local baseline dataset/camera/CUDA/batch/lr/num_workers/backbone settings +- [ ] Verify PID, GPU allocation, log tail, and SwanLab run URL. diff --git a/docs/superpowers/plans/2026-04-04-imf-horizon-grid-and-attnres-ablation.md b/docs/superpowers/plans/2026-04-04-imf-horizon-grid-and-attnres-ablation.md new file mode 100644 index 0000000..69e088d --- /dev/null +++ b/docs/superpowers/plans/2026-04-04-imf-horizon-grid-and-attnres-ablation.md @@ -0,0 +1,68 @@ +# IMF Horizon Grid and AttnRes Ablation Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Run a 6-run Phase-1 IMF horizon/action-step experiment grid across available GPUs, monitor progress and collect best rollout metrics, then use the best horizon setting for a Phase-2 visual-attnres ablation. + +**Architecture:** Use the current IMF training code as-is for Phase-1 by sweeping explicit `(pred_horizon, num_action_steps)` overrides while keeping emb=384, layer=12, and max_steps=50k fixed. Maintain a local experiment suite directory with a manifest and machine-readable status snapshots so progress can be resumed and summarized across turns. After Phase-1 completes, compare the current head-only attnres setup against a variant that also adds attnres into the visual ResNet path. + +**Tech Stack:** Python, Hydra/OmegaConf, PyTorch, SSH/Tailscale, JSON/CSV status files, SwanLab. + +--- + +### Task 1: Prepare the experiment suite manifest and state tracking + +**Files:** +- Create: `experiment_suites/2026-04-04-imf-horizon-grid/manifest.json` +- Create: `experiment_suites/2026-04-04-imf-horizon-grid/status.json` +- Create: `experiment_suites/2026-04-04-imf-horizon-grid/notes.md` + +- [ ] Define the 6 legal Phase-1 combinations: `(8,8)`, `(16,8)`, `(16,16)`, `(32,8)`, `(32,16)`, `(32,32)`. +- [ ] Record for each run: name, host, GPU slot, command, log path, SwanLab run name, and completion criteria. +- [ ] Define the comparison metric as the maximum rollout average reward seen during training (`max avg_reward`), preferably read from the best-checkpoint metadata and cross-checked against logs. +- [ ] Keep `status.json` updated with per-run state: queued / running / finished / failed plus latest parsed progress. + +### Task 2: Prepare the remote 8-GPU execution target + +**Files:** +- Remote working directory under `/home/droid/` +- Reuse or create a synced code directory for this suite + +- [ ] Verify the remote dataset path and environment path. +- [ ] Verify GPU availability and reserve 6 GPUs for Phase-1 launches. +- [ ] Sync the required code to a dedicated remote suite directory. +- [ ] Record exact remote paths back into the local suite manifest. + +### Task 3: Launch the 6 Phase-1 experiments in parallel + +**Files:** +- Reuse: `roboimi/demos/vla_scripts/train_vla.py` +- Modify only local suite tracking files unless a launch bug is discovered + +- [ ] Launch 6 runs concurrently with fixed settings: IMF, emb=384, layer=12, max_steps=50k. +- [ ] Keep all other relevant training hyperparameters aligned to the current strong baseline unless a concrete blocker appears. +- [ ] Assign one GPU per run on the 8xL20 host. +- [ ] Capture PID, log path, and SwanLab URL for each run in `status.json`. + +### Task 4: Monitor and summarize Phase-1 until all 6 finish + +**Files:** +- Update: `experiment_suites/2026-04-04-imf-horizon-grid/status.json` +- Update: `experiment_suites/2026-04-04-imf-horizon-grid/notes.md` + +- [ ] Periodically parse each run’s log/checkpoints to extract latest step, latest rollout reward, and best rollout reward so far. +- [ ] Keep a resumable local summary so progress can be continued in later turns without rediscovery. +- [ ] After all 6 runs finish, rank them by `max avg_reward` and write a compact Phase-1 summary. + +### Task 5: Prepare the Phase-2 visual-attnres ablation + +**Files:** +- Likely modify: vision backbone implementation and config files (to be confirmed after code inspection) +- Add/update targeted tests for the visual backbone path if code changes are needed + +- [ ] Use the best Phase-1 `(pred_horizon, num_action_steps)` combination as the fixed rollout setting for Phase-2. +- [ ] Compare: + 1. current setup: attnres only in the IMF head + 2. ablation setup: attnres in both IMF head and visual encoder path +- [ ] Keep the rest of the training settings fixed. +- [ ] Launch and monitor the Phase-2 pair after Phase-1 summary is complete. diff --git a/docs/superpowers/plans/2026-04-05-lewm-vit-backbone-implementation.md b/docs/superpowers/plans/2026-04-05-lewm-vit-backbone-implementation.md new file mode 100644 index 0000000..641fc46 --- /dev/null +++ b/docs/superpowers/plans/2026-04-05-lewm-vit-backbone-implementation.md @@ -0,0 +1,92 @@ +# LEWM ViT Backbone Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace the current ResNet visual encoder in roboimi VLA training with a frozen LEWM ViT visual backbone (encoder + projector) that consumes the three camera views jointly and outputs one 192-d CLS embedding per timestep, then launch two 50k runs on the 5880 machine. + +**Architecture:** Add a new joint-multiview LEWM backbone that fuses `front/top/r_vis` into one LEWM-style image, reproduces LEWM preprocessing, loads frozen weights from the trained checkpoint, and exposes a `joint_output_dim=192`. Add a minimal `VLAAgent` compatibility branch so conditions can be sized from joint visual dim instead of `output_dim * num_cams`, while leaving the rest of the diffusion pipeline unchanged. + +**Tech Stack:** PyTorch, transformers `ViTModel`, Hydra configs, existing roboimi VLA training/eval scripts, remote SSH/rsync to 100.73.14.65. + +--- + +### Task 1: Add failing tests for LEWM joint-vision backbone contract + +**Files:** +- Create: `tests/test_lewm_vit_backbone.py` +- Modify: `tests/test_imf_vla_agent.py` + +- [ ] **Step 1: Write the failing backbone shape/load test** +- [ ] **Step 2: Run `pytest tests/test_lewm_vit_backbone.py -q` and verify it fails** +- [ ] **Step 3: Extend `tests/test_imf_vla_agent.py` with a failing joint-output backbone case** +- [ ] **Step 4: Run `pytest tests/test_imf_vla_agent.py -q` and verify it fails** + +### Task 2: Implement LEWM joint-multiview frozen backbone + +**Files:** +- Create: `roboimi/vla/models/backbones/lewm_vit_backbone.py` +- Modify: `roboimi/vla/models/backbones/__init__.py` only if exports are needed + +- [ ] **Step 1: Create `LEWMViTBackbone` with public attrs `camera_names`, `num_cameras`, `joint_output_dim=192`** +- [ ] **Step 2: Reproduce LEWM preprocessing and joint multiview fusion** +- [ ] **Step 3: Load checkpoint weights from `model.encoder.*` and `model.projector.*`** +- [ ] **Step 4: Freeze encoder/projector and keep them in eval mode via `train()` override** +- [ ] **Step 5: Run `pytest tests/test_lewm_vit_backbone.py -q` and verify green** + +### Task 3: Add minimal agent support for joint visual dim + +**Files:** +- Modify: `roboimi/vla/agent.py` +- Test: `tests/test_imf_vla_agent.py` + +- [ ] **Step 1: Add a `joint_output_dim` branch in `VLAAgent.__init__` for `per_step_cond_dim` / `global_cond_dim`** +- [ ] **Step 2: Keep `_build_cond()` semantics unchanged except for matching the new dim contract** +- [ ] **Step 3: Run `pytest tests/test_imf_vla_agent.py -q` and verify green** + +### Task 4: Add Hydra configs for LEWM backbone training + +**Files:** +- Create: `roboimi/vla/conf/backbone/lewm_vit_diffusion.yaml` +- Create: `roboimi/vla/conf/agent/lewm_imf_attnres.yaml` + +- [ ] **Step 1: Add backbone config pointing to the new LEWM backbone** +- [ ] **Step 2: Add `agent=lewm_imf_attnres` config with 3 cameras and `head.cond_dim=208`** +- [ ] **Step 3: Verify Hydra instantiation with a one-shot compose smoke** + +### Task 5: Verify focused local tests + +**Files:** +- Reuse the above + +- [ ] **Step 1: Run `pytest tests/test_lewm_vit_backbone.py tests/test_imf_vla_agent.py tests/test_eval_vla_headless_import.py -q`** +- [ ] **Step 2: If needed, run one tiny local import/forward smoke** + +### Task 6: Sync to 5880 and remote smoke with real checkpoint + +**Files:** +- Remote target: `/home/droid/roboimi_suite_20260404` + +- [ ] **Step 1: Rsync modified source/config files to `100.73.14.65:/home/droid/roboimi_suite_20260404`** +- [ ] **Step 2: Run a 2-step smoke on GPU0 with `agent.head.n_emb=384`, `train.rollout_num_episodes=10`, real LEWM checkpoint** +- [ ] **Step 3: Run a 2-step smoke on GPU1 with `agent.head.n_emb=256`, same checkpoint** + +### Task 7: Launch two real 50k runs on the 5880 machine + +**Files:** +- Remote logs under `/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/` + +- [ ] **Step 1: Launch embed384/layer12 on GPU0** +- [ ] **Step 2: Launch embed256/layer12 on GPU1** +- [ ] **Step 3: Ensure both use `data.camera_names=[r_vis,top,front]`, `pred_horizon=16`, `num_action_steps=8`, `train.rollout_num_episodes=10`, `max_steps=50000`** +- [ ] **Step 4: Record run names, pids, log paths, SwanLab URLs** + +### Task 8: Update experiment tracking docs and commit + +**Files:** +- Create: `experiment_suites/2026-04-05-lewm-vit-transfer/manifest.json` +- Create: `experiment_suites/2026-04-05-lewm-vit-transfer/status.json` +- Create: `experiment_suites/2026-04-05-lewm-vit-transfer/notes.md` + +- [ ] **Step 1: Record checkpoint path, frozen LEWM design, rollout=10, and both run configs** +- [ ] **Step 2: Record running status after launch** +- [ ] **Step 3: Commit implementation + docs with a focused message** diff --git a/docs/superpowers/plans/2026-04-05-phase2-full-attnres-vision-plan.md b/docs/superpowers/plans/2026-04-05-phase2-full-attnres-vision-plan.md new file mode 100644 index 0000000..8414f81 --- /dev/null +++ b/docs/superpowers/plans/2026-04-05-phase2-full-attnres-vision-plan.md @@ -0,0 +1,64 @@ +# Phase-2 Full-AttnRes Vision Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace all ResNet residual units in the vision backbone with AttnRes-based image blocks while preserving the current IMF agent interfaces and launch a Phase-2 experiment anchored on the best Phase-1 horizon setting. + +**Architecture:** Keep the current multi-camera encoder shell and per-camera output contract, but introduce a new ResNet-like 2D AttnRes backbone that preserves stage-wise downsampling and final SpatialSoftmax conditioning. Wire it into the existing `ResNetDiffusionBackbone` via an opt-in mode and keep the agent/head/data interfaces unchanged. + +**Tech Stack:** PyTorch, Hydra/OmegaConf, existing IMF AttnRes transformer components, pytest. + +--- + +### Task 1: Add failing tests for the new full-AttnRes visual backbone + +**Files:** +- Create: `tests/test_attnres_resnet2d_backbone.py` +- Update: `tests/test_imf_vla_agent.py` + +- [ ] **Step 1: Write a failing backbone shape test** +- [ ] **Step 2: Run it to confirm the new backbone/config does not exist yet** +- [ ] **Step 3: Add a failing IMF agent wiring test for unchanged cond_dim=208** +- [ ] **Step 4: Run the targeted tests and capture the failure** + +### Task 2: Implement a ResNet-like 2D AttnRes backbone + +**Files:** +- Create: `roboimi/vla/models/backbones/attnres_resnet2d.py` +- Modify: `roboimi/vla/models/backbones/resnet_diffusion.py` + +- [ ] **Step 1: Add minimal 2D tokenization helpers and positional encoding / bias handling** +- [ ] **Step 2: Implement `AttnResImageBlock2D` for feature maps** +- [ ] **Step 3: Implement `AttnResResNetLikeBackbone2D` with stage-wise downsampling** +- [ ] **Step 4: Wire `_SingleRgbEncoder` to choose between original ResNet trunk and the new full-AttnRes trunk** +- [ ] **Step 5: Run the new backbone tests** + +### Task 3: Expose config switches and agent wiring + +**Files:** +- Modify: `roboimi/vla/conf/backbone/resnet_diffusion.yaml` +- Modify: `roboimi/vla/conf/agent/resnet_imf_attnres.yaml` + +- [ ] **Step 1: Add a backbone mode/config flag for the full-AttnRes vision trunk** +- [ ] **Step 2: Add defaults for attnres image depth/heads/etc. if needed** +- [ ] **Step 3: Add a Phase-2 launch override path that enables the new visual trunk** +- [ ] **Step 4: Run agent wiring tests again** + +### Task 4: Smoke-verify training path + +**Files:** +- Reuse existing training scripts and configs + +- [ ] **Step 1: Run a short CPU or tiny-step smoke instantiation / `compute_loss` test** +- [ ] **Step 2: If needed, run a very short training smoke launch** +- [ ] **Step 3: Verify no cond-dim or rollout-loading regressions** + +### Task 5: Launch the Phase-2 experiment + +**Files:** +- Update experiment tracking under `experiment_suites/` + +- [ ] **Step 1: Use Phase-1 best setting (`pred_horizon=16`, `num_action_steps=8`)** +- [ ] **Step 2: Launch baseline reference or reuse existing result** +- [ ] **Step 3: Launch full-AttnRes vision experiment** +- [ ] **Step 4: Track rollout metrics and compare max avg_reward** diff --git a/docs/superpowers/plans/2026-04-06-resnet-multitoken-imf.md b/docs/superpowers/plans/2026-04-06-resnet-multitoken-imf.md new file mode 100644 index 0000000..23e20fa --- /dev/null +++ b/docs/superpowers/plans/2026-04-06-resnet-multitoken-imf.md @@ -0,0 +1,81 @@ +# ResNet Multitoken IMF Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Implement a standard-ResNet-18 multiview IMF variant that emits three condition tokens per obs step and launch four L20 experiments for `n_emb in {256,384}` and `n_layer in {12,16}`. + +**Architecture:** The ResNet backbone will optionally return one token per camera instead of concatenating all cameras into one token. `VLAAgent` will pair each camera token with the current state, project each pair into a condition token, flatten the per-step camera tokens into one cond sequence, and feed that sequence into the existing IMF/AttnRes head. + +**Tech Stack:** PyTorch, torchvision ResNet-18, Hydra, pytest, SwanLab, SSH/Tailscale. + +--- + +### Task 1: Add failing tests for multi-token conditioning + +**Files:** +- Modify: `tests/test_imf_vla_agent.py` +- Modify: `tests/test_resnet_transformer_agent_wiring.py` + +- [ ] **Step 1: Add a direct agent test** + - Stub a vision backbone returning `(B,T,3,D)` and assert `_build_cond()` yields `(B, T*3, D_cond)`. + - Assert state is paired with each camera token, not concatenated across cameras first. +- [ ] **Step 2: Add Hydra wiring test** + - Instantiate a new `agent=resnet_imf_attnres_multitoken` config with small dims. + - Assert `condition_tokens_per_step == 3`, `condition_sequence_length == obs_horizon * 3`, and head `n_obs_steps` receives that sequence length. +- [ ] **Step 3: Run focused tests and verify RED** + - `python -m pytest tests/test_imf_vla_agent.py tests/test_resnet_transformer_agent_wiring.py -q` + +### Task 2: Implement multi-token ResNet conditioning path + +**Files:** +- Modify: `roboimi/vla/models/backbones/resnet_diffusion.py` +- Modify: `roboimi/vla/agent.py` +- Create: `roboimi/vla/conf/agent/resnet_imf_attnres_multitoken.yaml` + +- [ ] **Step 1: Extend ResNet backbone** + - Add an opt-in flag to return `(B,T,num_cams,D)` camera tokens instead of one concatenated `(B,T,num_cams*D)` token. + - Keep standard ResNet-18 vision mode; do not switch to AttnRes vision. +- [ ] **Step 2: Extend VLAAgent condition building** + - Support visual features with rank 4 `(B,T,K,D)`. + - Broadcast state to `(B,T,K,D_state)`, concatenate per camera, apply projector per token, then flatten to `(B,T*K,D_cond)`. + - Track `condition_tokens_per_step` and `condition_sequence_length`. +- [ ] **Step 3: Update transformer-head instantiation** + - Pass `n_obs_steps=condition_sequence_length` when building transformer heads. +- [ ] **Step 4: Add Hydra config** + - New agent config uses: + - separate ResNet-18 per camera + - standard residual vision trunk (`vision_backbone_mode=resnet`) + - condition projector output dim tied to `${agent.head.n_emb}` + - rollout episodes `10`, `pred_horizon=16`, `num_action_steps=8` + +### Task 3: Verify locally + +**Files:** +- Modify only if verification reveals issues + +- [ ] **Step 1: Run focused tests and make them pass** + - `python -m pytest tests/test_imf_vla_agent.py tests/test_resnet_transformer_agent_wiring.py -q` +- [ ] **Step 2: Run regression subset** + - `python -m pytest tests/test_eval_vla_headless.py tests/test_train_vla_rollout_validation.py tests/test_simple_robot_dataset_image_loading.py -q` +- [ ] **Step 3: Run local smoke instantiation** + - instantiate the new Hydra config and verify cond shape / sequence length + +### Task 4: Launch 4 L20 experiments + +**Files:** +- Remote repo copy under `/home/droid/roboimi_suite_20260404` + +- [ ] **Step 1: Sync code to `100.119.99.14`** +- [ ] **Step 2: Smoke the new config on remote** +- [ ] **Step 3: Launch runs** + - `(n_emb=256, n_layer=12)` + - `(n_emb=256, n_layer=16)` + - `(n_emb=384, n_layer=12)` + - `(n_emb=384, n_layer=16)` +- [ ] **Step 4: Keep fixed across runs** + - rollout episodes `10` + - `pred_horizon=16` + - `num_action_steps=8` + - standard ResNet-18 vision trunk + - three separate camera weights +- [ ] **Step 5: Record PIDs, GPUs, log paths, SwanLab URLs** diff --git a/docs/superpowers/plans/2026-04-06-siglip2-multiview-vla.md b/docs/superpowers/plans/2026-04-06-siglip2-multiview-vla.md new file mode 100644 index 0000000..541555e --- /dev/null +++ b/docs/superpowers/plans/2026-04-06-siglip2-multiview-vla.md @@ -0,0 +1,78 @@ +# SigLIP2 Multiview VLA Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Integrate a frozen shared SigLIP2 multiview encoder into the IMF/AttnRes policy, preserve raw-256 image handling, and launch two 50k-step experiments on the 5880 host with per-view projection dims 96 and 192. + +**Architecture:** A new backbone will independently encode each camera view with SigLIP2 and project each 768-d pooled feature to a configurable per-view dimension. `VLAAgent` will concatenate visual features with robot state, then optionally project the combined per-step condition to the head's required 384-d interface before diffusion training/inference. + +**Tech Stack:** PyTorch, transformers SigLIP2, Hydra, pytest, SSH/Tailscale, SwanLab. + +--- + +### Task 1: Add failing tests for SigLIP2 backbone and projected conditioning + +**Files:** +- Create: `tests/test_siglip2_diffusion_backbone.py` +- Modify: `tests/test_imf_vla_agent.py` + +- [ ] **Step 1: Write failing backbone tests** + - Instantiate the new backbone with a stub SigLIP2 vision model. + - Assert raw dataset resize is `None`, eval resize is `(256, 256)`, output shape is `(B, T, 3 * per_view_output_dim)`. + - Assert three views are encoded independently and projected. +- [ ] **Step 2: Run focused tests and verify RED** + - Run `pytest tests/test_siglip2_diffusion_backbone.py tests/test_imf_vla_agent.py -q` + - Expect failure because the backbone/config/projector do not exist yet. +- [ ] **Step 3: Extend agent wiring tests** + - Add a Hydra/instantiate test for a new SigLIP2 IMF config. + - Assert raw condition dim `3 * per_view_output_dim + obs_dim`, projected cond dim `384`, and head `cond_dim == 384`. + +### Task 2: Implement SigLIP2 backbone and optional condition projector + +**Files:** +- Create: `roboimi/vla/models/backbones/siglip2_diffusion_backbone.py` +- Create: `roboimi/vla/conf/backbone/siglip2_diffusion.yaml` +- Create: `roboimi/vla/conf/agent/siglip2_imf_attnres.yaml` +- Create: `roboimi/vla/conf/modules/linear_condition_projector.yaml` +- Modify: `roboimi/vla/models/backbones/__init__.py` +- Modify: `roboimi/vla/agent.py` + +- [ ] **Step 1: Implement backbone** + - Load `SiglipVisionModel.from_pretrained("google/siglip2-base-patch16-256")`. + - Normalize `[0,1]` pixels with mean/std `0.5` and encode each view independently. + - Project each 768-d pooled feature to configurable per-view dim and concatenate across cameras. +- [ ] **Step 2: Implement optional condition projector** + - Allow `VLAAgent` to accept `cond_projector`. + - Track `raw_per_step_cond_dim` and projected `per_step_cond_dim` / `global_cond_dim`. + - Apply the projector in `_build_cond()` after visual+state concatenation. +- [ ] **Step 3: Add Hydra configs** + - New agent config should default to `n_emb=384`, `n_layer=12`, `pred_horizon=16`, `num_action_steps=8`, `head.cond_dim=384`. + - Backbone config should set `dataset_image_resize_shape: null` and `eval_image_resize_shape: [256, 256]`. + +### Task 3: Verify locally and prepare remote execution + +**Files:** +- Modify as needed only if tests/smoke reveal issues + +- [ ] **Step 1: Run focused tests and make them pass** + - `pytest tests/test_siglip2_diffusion_backbone.py tests/test_imf_vla_agent.py tests/test_eval_vla_headless.py tests/test_train_vla_rollout_validation.py tests/test_simple_robot_dataset_image_loading.py -q` +- [ ] **Step 2: Run a local smoke instantiation** + - Instantiate the new Hydra config with stubbed optional modules or offline-safe monkeypatching. +- [ ] **Step 3: Review diffs for unintended LEWM/raw256 regressions** + +### Task 4: Sync to 5880 and launch experiments + +**Files:** +- Remote repo copy under `/home/droid/roboimi_suite_20260404` + +- [ ] **Step 1: Stop superseded remote jobs** +- [ ] **Step 2: Sync updated code to remote** + - Prefer `rsync` or `git push/pull` without overwriting unrelated files. +- [ ] **Step 3: Remote smoke test** + - Confirm SigLIP2 model download/import works in `/home/droid/miniforge3/envs/roboimi/bin/python`. + - Confirm headless rollout path still uses `256x256` eval resize. +- [ ] **Step 4: Launch experiment A** + - `per_view_output_dim=96`, `embed=384`, `layer=12`, `pred=16`, `exec=8`, `steps=50000`. +- [ ] **Step 5: Launch experiment B** + - `per_view_output_dim=192`, same other hyperparameters. +- [ ] **Step 6: Record PIDs, GPUs, log paths, and SwanLab run URLs.** diff --git a/docs/superpowers/plans/2026-04-23-sim-air-insert-ring-bar.md b/docs/superpowers/plans/2026-04-23-sim-air-insert-ring-bar.md new file mode 100644 index 0000000..54c7f28 --- /dev/null +++ b/docs/superpowers/plans/2026-04-23-sim-air-insert-ring-bar.md @@ -0,0 +1,311 @@ +# sim_air_insert_ring_bar Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add an independent dual-Diana MuJoCo task `sim_air_insert_ring_bar` with a square ring block, a square bar block, staged rewards, strict finite-geometry in-air insertion success detection, and a task-specific scripted policy. + +**Architecture:** Reuse the current dual-Diana EE-control stack and environment factory, but add a task-specific scene XML, robot asset entrypoint, sampling helpers, and a new task-specific environment module. Keep `sim_transfer` untouched while introducing pure-Python geometry helpers and focused tests so reward/success behavior can be regression tested without requiring a full MuJoCo rollout in every test. + +**Tech Stack:** Python, unittest, MuJoCo XML assets, existing dual-Diana environment classes, Hydra-compatible task naming/config patterns. + +--- + +## File Structure / Responsibilities + +- **Create:** `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml` + - Defines the rigid ring body and bar body, each with a free joint and stable box-based geoms. +- **Create:** `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml` + - Scene entrypoint that includes the shared world/table/robot assets plus the new object XML. +- **Modify:** `roboimi/assets/robots/diana_med.py` + - Add a task-specific robot asset class for the new scene XML without changing existing `BiDianaMed` behavior. +- **Modify:** `roboimi/utils/act_ex_utils.py` + - Add deterministic helpers to sample left/right planar placement regions for ring and bar objects. +- **Modify:** `roboimi/utils/constants.py` + - Register the new task name and default metadata. +- **Create:** `roboimi/envs/double_air_insert_env.py` + - New task-specific environment, finite-geometry success helpers, reset logic, reward logic, and task factory branch. +- **Modify:** `roboimi/envs/double_pos_ctrl_env.py` + - Route `make_sim_env()` to the new task-specific environment while keeping current `sim_transfer` logic unchanged. +- **Create:** `roboimi/demos/diana_air_insert_policy.py` + - Task-specific waypoint/open-loop scripted policy for grasp-lift-align-insert. +- **Modify:** `roboimi/demos/vla_scripts/eval_vla.py` + - Reset the new task with the correct sampled task state instead of assuming a single transfer box pose. +- **Create:** `tests/test_air_insert_env.py` + - Focused unit tests for sampling, reset helpers, reward progression, and strict success detection. +- **Modify:** `tests/test_eval_vla_headless.py` + - Add coverage that headless evaluation dispatches the correct reset sampler for the new task. +- **Modify:** `tests/test_robot_asset_paths.py` + - Verify the new robot asset class resolves its XML path correctly independent of cwd. + +--- + +### Task 1: Add failing tests for task registration, samplers, and asset wiring + +**Files:** +- Create: `tests/test_air_insert_env.py` +- Modify: `tests/test_eval_vla_headless.py` +- Modify: `tests/test_robot_asset_paths.py` +- Modify: `roboimi/utils/act_ex_utils.py` (later in implementation) +- Modify: `roboimi/utils/constants.py` (later in implementation) +- Modify: `roboimi/assets/robots/diana_med.py` (later in implementation) +- Modify: `roboimi/envs/double_pos_ctrl_env.py` (later in implementation) +- Create: `roboimi/envs/double_air_insert_env.py` (minimal stub in this task) + +- [ ] **Step 1: Write failing tests for task config and sampling helpers** + +Add tests in `tests/test_air_insert_env.py` covering: +- `SIM_TASK_CONFIGS['sim_air_insert_ring_bar']` exists +- `sample_air_insert_ring_bar_pose()` (or equivalent helper) returns ring/bar positions with fixed z and correct left/right planar ranges +- output structure is explicit and easy for reset/eval code to consume + +- [ ] **Step 2: Write failing tests for environment factory dispatch and robot asset resolution** + +Add tests covering: +- `make_sim_env('sim_air_insert_ring_bar', headless=True)` dispatches to the new environment with rendering disabled +- a new robot asset class resolves the new XML path independent of cwd, similar to the existing `BiDianaMed` test pattern + +- [ ] **Step 3: Write failing tests for eval reset helper dispatch** + +Extend `tests/test_eval_vla_headless.py` so headless eval can reset the new task using the new sampler instead of hard-coding `sample_transfer_pose()`. + +- [ ] **Step 4: Run the targeted tests to verify they fail for the expected missing-feature reasons** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v` + +Expected: +- FAIL because the new task config/helper/class/dispatch branch does not exist yet + +- [ ] **Step 5: Implement the minimal production code to satisfy the new task registration and helper tests** + +Implement only enough to make the new tests pass: +- add new task config entry +- add the new placement sampler +- add the new robot asset class +- create a minimal importable `double_air_insert_env.py` stub and class/function surface needed for factory dispatch tests +- add the factory dispatch branch / headless wiring +- update eval reset dispatch for the new task + +- [ ] **Step 6: Re-run the targeted tests to verify they pass** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v` + +Expected: +- PASS for the new registration/sampler/dispatch/asset tests + +- [ ] **Step 7: Commit Task 1** + +Run: +`git add tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/utils/act_ex_utils.py roboimi/utils/constants.py roboimi/assets/robots/diana_med.py roboimi/envs/double_pos_ctrl_env.py roboimi/envs/double_air_insert_env.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(env): register sim air insert ring bar task"` + +--- + +### Task 2: Add the MuJoCo ring+bar scene assets and reset helpers + +**Files:** +- Create: `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml` +- Create: `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml` +- Create or Modify: `roboimi/envs/double_air_insert_env.py` +- Modify: `tests/test_air_insert_env.py` + +- [ ] **Step 1: Write failing tests for object reset helpers and scene-specific joint naming assumptions** + +In `tests/test_air_insert_env.py`, add unit tests for helper functions that: +- write ring pose to `ring_block_joint` from the named task-state mapping +- write bar pose to `bar_block_joint` from the named task-state mapping +- read back `env_state` as a stable 14D vector `[ring_pos, ring_quat, bar_pos, bar_quat]` + +Use fake `mj_data` objects so tests stay fast and deterministic. + +- [ ] **Step 2: Run the focused test slice and verify it fails** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v` + +Expected: +- FAIL because reset/state helper functions and joint conventions are not implemented yet + +- [ ] **Step 3: Implement the scene XML files and reset/state helper code** + +Implement: +- the object XML with one rigid ring body and one rigid bar body +- the task scene XML entrypoint using the shared world/table/robot includes +- reset helper(s) in `double_air_insert_env.py` that set qpos for both free joints with fixed quaternions +- task-state accessor(s) returning both object poses in a stable structure + +- [ ] **Step 4: Re-run the focused test slice and verify it passes** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v` + +Expected: +- PASS for reset/state helper tests + +- [ ] **Step 5: Commit Task 2** + +Run: +`git add roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(scene): add ring and bar insertion scene assets"` + +--- + +### Task 3: Implement strict reward and finite-geometry success detection + +**Files:** +- Modify: `roboimi/envs/double_air_insert_env.py` +- Modify: `tests/test_air_insert_env.py` + +- [ ] **Step 1: Write failing tests for reward stages and strict success detection** + +Add tests in `tests/test_air_insert_env.py` for: +- left contact stage reward +- right contact stage reward +- ring lifted off table stage +- bar lifted off table stage +- positive success case where a finite bar truly passes through the aperture +- negative case where the centerline would pass but the finite square body would clip +- negative case where the bar has not crossed the ring thickness direction enough +- negative case where one/both objects are still on the table + +Structure the tests around pure helper functions and light fake contact/state objects so the geometry logic is directly regression tested. + +- [ ] **Step 2: Run the focused tests and verify they fail for missing reward/success logic** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v` + +Expected: +- FAIL because the staged reward and finite-geometry insertion logic are not implemented yet + +- [ ] **Step 3: Implement minimal strict success helpers and reward logic** + +Implement in `roboimi/envs/double_air_insert_env.py`: +- pure helper(s) for transforming bar geometry into ring-local coordinates +- finite-geometry insertion predicate (not centerline-only) +- table-contact / airborne checks +- staged reward function returning the highest achieved stage with `max_reward = 5` + +- [ ] **Step 4: Re-run the focused tests to verify the logic passes** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v` + +Expected: +- PASS for reward and success-detection regression tests + +- [ ] **Step 5: Commit Task 3** + +Run: +`git add roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(env): add strict air insertion reward and success logic"` + +--- + +### Task 4: Add the scripted policy and integration smoke coverage + +**Files:** +- Create: `roboimi/demos/diana_air_insert_policy.py` +- Modify: `roboimi/demos/diana_record_sim_episodes.py` +- Modify: `tests/test_air_insert_env.py` +- Optionally Modify: `roboimi/demos/vla_scripts/eval_vla.py` (only if integration gaps remain after Task 1) + +- [ ] **Step 1: Write failing tests for scripted-policy action shape and basic generation** + +Add tests covering: +- the new policy produces a 16D action +- trajectory generation accepts sampled named task state without error +- the first action is a valid open-gripper safe pose command +- a deterministic nominal smoke path (with canonical sampled state or fake env shim) reaches the intended terminal interface contract without shape/reward mismatches + +Keep the tests unit-level; do not require a full MuJoCo rollout for every assertion. + +- [ ] **Step 2: Write failing tests for the scripted rollout entrypoint and a real headless smoke path** + +Add coverage for both: +- the standard scripted rollout entrypoint (`roboimi/demos/diana_record_sim_episodes.py`) can select the new task sampler/policy instead of remaining sim_transfer-only +- a deterministic integration/smoke test that instantiates `make_sim_env('sim_air_insert_ring_bar', headless=True)`, resets with sampled named task state, and steps a few actions or scripted-policy outputs using the real task XML and task-specific wiring + +- [ ] **Step 3: Run the scripted-policy tests and verify they fail** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v` + +Expected: +- FAIL because the new scripted policy does not exist yet + +- [ ] **Step 4: Implement the waypoint-based scripted policy** + +Implement a conservative open-loop policy with phases: +- safe wait pose +- above-target approach +- descend + grasp +- dual lift +- airborne meeting alignment +- bar push-through insertion + +Use fixed orientations for version 1 and follow the existing repository style from `diana_policy.py`. + +- [ ] **Step 5: Re-run the scripted-policy tests to verify they pass** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v` + +Expected: +- PASS for scripted-policy tests + +- [ ] **Step 6: Run the combined verification suite for this feature** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_eval_vla_rollout_artifacts tests.test_train_vla_rollout_validation tests.test_robot_asset_paths -v` + +Expected: +- PASS with 0 failures + +- [ ] **Step 6b: Run the mandatory real headless smoke check** + +Run a focused smoke command that instantiates the real task, resets with sampled state, and steps a few actions using the new scripted policy or a deterministic action sequence. + +Example command (adjust module/test helper if needed): +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env.AirInsertEnvSmokeTest -v` + +Expected: +- PASS, proving the real XML/assets/env wiring instantiate and step correctly in headless mode + +- [ ] **Step 7: Commit Task 4** + +Run: +`git add roboimi/demos/diana_air_insert_policy.py tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(policy): add scripted air insertion policy"` + +--- + +### Task 5: Final verification and implementation review + +**Files:** +- Review all files touched above + +- [ ] **Step 1: Run fresh end-to-end verification before claiming completion** + +Run: +`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v` + +Expected: +- PASS with 0 failures + +- [ ] **Step 2: Inspect git status and recent commits** + +Run: +`git status --short && git log --oneline --decorate -n 8` + +Expected: +- only intended feature files modified / committed + +- [ ] **Step 3: Request final code review for the completed feature** + +Use the requesting-code-review skill against the full diff from the feature branch starting point to current HEAD. + +- [ ] **Step 4: Address any review findings and re-run verification if code changes** + +If fixes are made, repeat the unittest command from Step 1. + +- [ ] **Step 5: Hand off using finishing-a-development-branch** + +After verification and review, use the finishing-a-development-branch skill to decide merge / PR / cleanup. diff --git a/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md b/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md new file mode 100644 index 0000000..f232125 --- /dev/null +++ b/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md @@ -0,0 +1,241 @@ +# VLA Training + Headless Rollout + SwanLab Design + +**Date:** 2026-03-30 +**Branch:** feat-align-dp-transformer-ee + +## Goal +在当前仓库中补齐默认 `resnet_transformer` / `Transformer1D` 路线的训练依赖，使用数据集 `/home/droid/project/diana_sim/sim_transfer` 启动训练；同时支持训练过程中的 SwanLab 标量日志上传，并为后续 rollout 验证提供 headless 模式，避免弹出 MuJoCo / OpenCV 图形界面。 + +## Non-Goals +- 不重写整套训练框架 +- 不引入新的 workspace / callback 框架 +- 不在本轮做复杂的视频/媒体日志上传 +- 不修改数据集格式本身 + +## Current State +- 默认训练配置已切到 `agent=resnet_transformer`，head 为 `Transformer1D` +- 当前环境缺少训练所需的若干 Python 依赖：`diffusers`、`torchvision`、`einops`、`swanlab` +- 评估环境 `make_sim_env(task_name)` 当前写死 `is_render=True` +- 相机线程 `camera_viewer()` 默认会 `cv2.namedWindow/imshow`，即使只想拿图像也会弹窗 +- 训练脚本当前支持 train/val loss、checkpoint，但没有 SwanLab 集成 +- 数据集目录 `/home/droid/project/diana_sim/sim_transfer` 下已有 100 个 episode，但还没有 `dataset_stats.pkl` + +## User Requirements +1. 在现有 mamba 环境里补齐训练依赖 +2. 在 `/home/droid/project/diana_sim/sim_transfer` 上开始训练 +3. 如果训练中需要 rollout 验证，希望支持 headless，不弹 GUI +4. 训练指标上传到 SwanLab +5. 默认 SwanLab project 名为 `roboimi-vla` + +## Proposed Approach +采用“最小必要改造”方案： + +### 1. Dependency Layer +在现有 `roboimi` 环境中补齐缺失训练依赖，并优先保持现有环境名与脚本入口不变。 + +#### Install Plan +- 环境：继续使用现有 mamba 环境 `roboimi` +- 安装方式： + - 优先使用当前 env 的 `python -m pip install` + - 安装包： + - `diffusers` + - `torchvision` + - `einops` + - `swanlab` +- 版本策略： + - 优先选择与当前 `torch==2.4.0` 可兼容的最新可安装版本 + - 若出现兼容性问题，再回退到与 `torch 2.4` 对齐的稳定版本 +- 复现策略： + - 本轮会把**实际安装成功的 resolved versions** 补写回仓库的环境定义文件，避免后续环境漂移 + +训练前验证以下 import： +- `torch` +- `hydra` +- `omegaconf` +- `diffusers` +- `torchvision` +- `einops` +- `swanlab` +- `cv2` +- `h5py` +- `mujoco` + +### 2. Dataset Preparation +直接复用现有 `SimpleRobotDataset`，仅将 `data.dataset_dir` 指向： +- `/home/droid/project/diana_sim/sim_transfer` + +训练前使用现有统计脚本生成： +- `/home/droid/project/diana_sim/sim_transfer/dataset_stats.pkl` + +统计文件生成命令目标为： +- 从仓库根目录执行 +- 直接针对 `/home/droid/project/diana_sim/sim_transfer` 输出 stats +- 训练脚本不再依赖默认数据目录 + +### 3. SwanLab Logging +在训练脚本中增加一个轻量 logging 集成层： +- 通过配置决定是否启用 SwanLab，默认启用 +- 默认 project：`roboimi-vla` +- API key 不写入仓库，不写入配置文件，只通过本地登录状态或环境变量使用 +- 当 `train.use_swanlab=true` 时： + - 若 `swanlab` 不可 import，训练直接 fail fast + - 若未登录或认证失败，训练直接 fail fast +- 每个训练日志点上传： + - `train/loss` + - `train/lr` + - `train/best_loss` + - `train/step` +- 每次验证时上传： + - `val/loss` +- 训练结束时记录最终 checkpoint 路径与 best checkpoint 路径 + +### 4. Headless Rollout Design +目标是让 rollout 验证可以“拿到图像观测，但不弹任何窗口”。 + +最小改造策略： +- 给 `make_sim_env(...)` 增加 `headless` / `is_render` 参数 +- 给相机线程显示逻辑增加开关： + - headless 时继续更新 `r_vis/top/front/...` 图像缓存 + - 但不执行 `cv2.namedWindow` / `cv2.imshow` / `cv2.waitKey` +- 评估脚本中： + - headless 时不调用 `env.render()` + - 仍然允许 `env._get_image_obs()` 和 policy inference 正常运行 + +#### Training-Time Rollout Scope +- 本轮**会提供一个可选的 checkpoint-time rollout validation 路径**，默认关闭 +- 启用后，在训练保存 checkpoint 时可以调用同仓库的 rollout/eval 逻辑做少量 episode 验证 +- 此路径要求支持**唯一权威开关** `eval.headless=true`，即： + - 不弹 MuJoCo viewer + - 不执行 `cv2.namedWindow / cv2.imshow / cv2.waitKey` + - 仍可读取图像并完成策略推理 +- 默认情况下不增加频繁 rollout，以避免拖慢训练；只提供能力与配置开关 + +如果验证发现相机线程强依赖 GUI，我们的降级策略是： +- 训练主流程 + SwanLab 必须先跑通 +- rollout validation 保持为显式可选能力 +- 但本轮仍要保证至少存在可调用的 headless 验证执行路径，而不是仅停留在文档层面 + +### 5. Training Execution Strategy +分两步执行： + +#### Step A: Smoke Run +使用较小步数启动一次 smoke training，确认： +- 数据集可正常读取 +- 统计文件可加载 +- 模型可实例化 +- 单步前后向正常 +- checkpoint 正常写出 +- SwanLab 成功上传标量 + +#### Step B: Real Training Run +在 smoke run 成功后，再启动正式训练。 + +## Execution Commands + +### A. Stats Generation +从仓库根目录执行，生成： +- `/home/droid/project/diana_sim/sim_transfer/dataset_stats.pkl` + +命令模板： +```bash +/home/droid/.conda/envs/roboimi/bin/python roboimi/vla/scripts/calculate_stats.py \ + --dataset_dir /home/droid/project/diana_sim/sim_transfer +``` + +### B. Smoke Training Command +从仓库根目录执行，核心覆盖项包括： +- `data.dataset_dir=/home/droid/project/diana_sim/sim_transfer` +- 较小 `train.max_steps` +- 较高日志频率 +- 启用 SwanLab +- 输出目录使用当前运行目录下的 `checkpoints/` + +命令模板： +```bash +/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \ + data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \ + train.max_steps=20 \ + train.log_freq=1 \ + train.save_freq=10 \ + train.use_swanlab=true \ + train.swanlab_project=roboimi-vla \ + train.rollout_validate_on_checkpoint=false +``` + +### C. Real Training Command +从仓库根目录执行，核心覆盖项包括： +- `data.dataset_dir=/home/droid/project/diana_sim/sim_transfer` +- 正式 `train.max_steps` +- 默认 project=`roboimi-vla` +- 若启用 rollout validation，则传入 `eval.headless=true` 以及训练侧 rollout 开关 + +命令模板： +```bash +/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \ + data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \ + train.use_swanlab=true \ + train.swanlab_project=roboimi-vla \ + train.rollout_validate_on_checkpoint=true \ + eval.headless=true +``` + +### D. Output Behavior +- checkpoint 输出目录：当前工作目录下的 `checkpoints/` +- 关键文件： + - `checkpoints/vla_model_step_.pt` + - `checkpoints/vla_model_best.pt` + - `checkpoints/vla_model_final.pt` + +## File-Level Changes +- `environment.yml` + - 补写新增训练依赖，保证后续可复现 +- `roboimi/demos/vla_scripts/train_vla.py` + - 增加 SwanLab 集成 + - 增加更明确的数据集目录覆盖支持 + - 增加可选 checkpoint-time rollout validation 入口 + - 保持当前 optimizer 对齐逻辑不变 +- `roboimi/vla/conf/config.yaml` + - 增加/扩展训练日志、SwanLab、rollout 相关配置项 +- `roboimi/vla/conf/eval/eval.yaml` + - 增加 `headless` 等评估控制项 +- `roboimi/envs/double_pos_ctrl_env.py` + - `make_sim_env` 支持 headless / no-render +- `roboimi/envs/double_base.py` + - 相机采集与 GUI 显示解耦 +- `roboimi/vla/scripts/calculate_stats.py` + - 改为直接支持通过命令行传入外部 `dataset_dir` +- tests（新增） + - 覆盖 SwanLab 可选初始化路径 + - 覆盖 headless 环境下“不弹窗但可取图”的关键逻辑 + +## Validation Plan +1. 补齐依赖后验证 import 全通过 +2. 生成 `dataset_stats.pkl` +3. 运行训练 smoke run +4. 确认 SwanLab dashboard 在 project `roboimi-vla` 下有标量更新 +5. 若启用 rollout 验证：确认 headless 下不弹 GUI，且 rollout 路径能真正执行 +6. 再启动正式训练 + +## Config Contract +本轮新增/固定的配置键以以下形式为准： +- `train.use_swanlab: true|false` +- `train.swanlab_project: roboimi-vla` +- `train.rollout_validate_on_checkpoint: true|false` +- `eval.headless: true|false` + +## Risks and Mitigations +- **Risk:** GUI/相机线程与离屏渲染耦合 + - **Mitigation:** 先解耦显示与图像更新；必要时把 rollout 验证降级为第二阶段 +- **Risk:** 现有 env 依赖不完整 + - **Mitigation:** 先做 import 验证，再做 smoke run +- **Risk:** 数据集过大导致 smoke run 也很慢 + - **Mitigation:** smoke run 只跑极小步数 +- **Risk:** SwanLab API key 泄漏 + - **Mitigation:** 不写入代码/配置，只保存在本地登录态或环境变量 + +## Success Criteria +- 训练脚本能在 `/home/droid/project/diana_sim/sim_transfer` 上启动 +- 能成功写出 checkpoint 到 `checkpoints/` +- SwanLab 在 `roboimi-vla` 项目下能看到 train/val 标量 +- headless rollout 具备不弹 GUI 的执行路径 +- 若训练侧启用 rollout validation，则该路径可以在 headless 模式下被实际调用 diff --git a/docs/superpowers/specs/2026-03-31-rollout-artifacts-design.md b/docs/superpowers/specs/2026-03-31-rollout-artifacts-design.md new file mode 100644 index 0000000..1d30446 --- /dev/null +++ b/docs/superpowers/specs/2026-03-31-rollout-artifacts-design.md @@ -0,0 +1,16 @@ +# Rollout Artifacts Design + +**Goal:** Add a one-off evaluation path that can record rollout video, export per-step timing breakdowns, and save executed end-effector trajectories for a selected checkpoint while preserving default eval behavior when artifact capture is disabled. + +**Approach:** Extend `roboimi/demos/vla_scripts/eval_vla.py` with optional evaluation-time artifact capture that stays backward compatible when disabled. Reuse existing environment observation and camera streams, record one camera stream to MP4, collect per-step timing around observation read / preprocessing / model inference / env step / total loop, and save per-step raw predicted EE actions plus executed EE poses after stepping. + +**Artifact contract:** +- `video.mp4`: optional MP4 encoded from a selected camera stream (`r_vis`, `top`, `front`, etc.), written only when recording is enabled. +- `trajectory.npz`: canonical trajectory export containing at minimum `step`, `reward`, `raw_action`, `executed_left_link7_pos`, `executed_left_link7_quat`, `executed_right_link7_pos`, `executed_right_link7_quat`, and optional duplicated tool-body poses if captured. +- `timing.json`: JSON-serializable per-episode timing summary with millisecond units for `obs_read_ms`, `preprocess_ms`, `inference_ms`, `env_step_ms`, `loop_total_ms`, plus aggregate mean/std/min/max and counts. Raw per-step timing arrays should also be persisted in the NPZ for later analysis. + +**Checkpoint selection:** Prefer an explicitly requested checkpoint path. If the caller asks for “latest” or omits a path in the execution helper, select the newest fully written checkpoint file by mtime/name and fail clearly if none exists. + +**Stop-training / execution safety:** Before rollout, stop any active training process using the target run, wait for process exit, then verify the chosen checkpoint exists and is readable. If the most recent checkpoint is missing or mid-write, fall back to the previous completed checkpoint or `vla_model_best.pt` with the decision logged. + +**Backward compatibility:** With all new eval flags left at default values, `_run_eval` return shape must remain compatible with existing callers, training-time rollout validation should continue to work without passing new options, and no artifact files should be written. diff --git a/docs/superpowers/specs/2026-04-01-imf-attnres-policy-design.md b/docs/superpowers/specs/2026-04-01-imf-attnres-policy-design.md new file mode 100644 index 0000000..eea9e7b --- /dev/null +++ b/docs/superpowers/specs/2026-04-01-imf-attnres-policy-design.md @@ -0,0 +1,272 @@ +# IMF-AttnRes Policy Migration Design + +**Date:** 2026-04-01 +**Status:** Approved in chat, written spec pending review + +## Goal + +将 `/home/droid/project/diffusion_policy` 中提交 `185ed659` 的 IMF-AttnRes diffusion policy 迁移到当前 `roboimi` 仓库，作为当前 DiT / Transformer diffusion policy 的替代训练选项；同时迁移其训练目标与一步推理机制，并保持 RoboIMI 现有的仿真环境、三相机视觉输入、数据集格式、训练脚本和 rollout 验证工作流可继续使用。 + +## Non-Goals + +- 不迁移 external repo 中与当前任务无关的 obs encoder、dataset、env wrapper、PushT 专用逻辑。 +- 不强行复刻 external repo 中全部目录结构；仅迁移当前 RoboIMI 训练所必需的模型、loss、inference 语义。 +- 不在本次工作中同时保留旧 DiT 为默认训练目标；旧配置继续可用，但新模型单独提供 config 入口。 + +## User-Confirmed Requirements + +1. 迁移对象是 `185ed659` 中的 **IMF-AttnRes 模型相关代码**。 +2. 不只是迁移骨架，还要迁移： + - **训练目标** + - **一步推理机制** +3. 视觉输入与当前 RoboIMI diffusion policy 一致： + - 使用三个相机图像作为条件输入 + - 图像观测必须作为条件，而不是拼进输出预测目标 +4. 当前任务里，IMF policy 用来替代现有 DiT/Transformer diffusion policy 训练。 +5. 训练参数沿用最近一次训练的大体设置（后续由训练命令显式覆盖），但推理方式改为 IMF 的 one-step 机制。 +6. 用户接受 IMF 中“全注意力 / 非因果注意力”的实现约束。 + +## External Source of Truth + +迁移语义以 external repo 的以下文件为准： + +- `diffusion_policy/model/diffusion/attnres_transformer_components.py` +- `diffusion_policy/model/diffusion/imf_transformer_for_diffusion.py` +- `diffusion_policy/policy/imf_transformer_hybrid_image_policy.py` +- 参考配置：`image_pusht_diffusion_policy_dit_imf_attnres_full.yaml` + +其中最关键的差异是：该策略并非 DDPM/DDIM 多步去噪，而是 IMF 训练目标 + one-step 推理。 + +## Current RoboIMI Baseline + +当前 RoboIMI 中与该任务直接相关的基线如下： + +- 视觉编码：`ResNetDiffusionBackbone` + - 三相机：`r_vis`, `top`, `front` + - 每个时间步将相机特征与 `qpos` 拼接为 per-step condition +- 策略主体：`VLAAgent` + - `compute_loss()` 使用 DDPM 噪声预测损失 + - `predict_action()` 使用 DDIM 多步采样 + - 在线控制通过动作队列机制在 `select_action()` 中按 chunk 触发预测 +- 训练脚本：`roboimi/demos/vla_scripts/train_vla.py` + - 支持 GPU 训练、SwanLab 日志、headless rollout 验证 + +因此，本次迁移的核心不是换视觉 backbone，而是替换 **head + loss + inference semantics**。 + +## Recommended Integration Approach + +采用 **最小侵入式集成**： + +1. **保留当前 RoboIMI 的视觉编码、数据读取、rollout/eval、训练脚本主框架**。 +2. **新增 IMF 专用 head 模块**，在 RoboIMI 内本地实现： + - AttnRes 组件 + - IMF transformer 主体 +3. **新增 IMF 专用 agent**，复用当前 `VLAAgent` 的： + - 归一化逻辑 + - 相机顺序管理 + - 观测缓存 / 动作 chunk 缓存 + - rollout 接口 + 但覆盖： + - `compute_loss()` + - `predict_action()` +4. **新增独立 Hydra config**，让 IMF policy 作为新的 agent 选项，不破坏已有 resnet_transformer / gr00t_dit 配置。 + +这样做的原因： + +- 迁移 IMF 语义时不必把当前 DDPM agent 搅乱； +- rollout / eval / checkpoint 逻辑仍然可复用； +- 便于和现有 Transformer / DiT 直接做 A/B 对比训练。 + +## Architecture + +### 1. Observation / Conditioning Path + +沿用当前 RoboIMI 的视觉路径： + +- 输入观测：`images={r_vis, top, front}` + `qpos` +- `ResNetDiffusionBackbone` 对每个相机编码，得到 per-camera feature +- `state_encoder` 编码 `qpos` +- 将三相机特征与 state feature 按时间步拼接，形成 `per_step_cond` + +这里不迁移 external repo 的 obs_encoder 实现；我们只对齐 **“图像作为条件 token 输入 transformer”** 这一语义。 + +### 2. Condition Tokenization + +对齐 external IMF transformer 的 token 使用方式： + +- action trajectory token：由 `(B, pred_horizon, action_dim)` 通过线性层映射到 `n_emb` +- time token：两个标量 `r` 与 `t`，分别通过 sinusoidal embedding + linear projection 得到 token +- observation token：`per_step_cond` 通过线性层映射到 `n_emb` +- 最终 token 序列为： + - `[r_token, t_token, obs_cond_tokens..., action_tokens...]` + +在当前任务中，obs token 数量等于 `obs_horizon`，且图像观测始终作为条件输入。 + +### 3. IMF-AttnRes Backbone + +在 RoboIMI 内新增 AttnRes backbone 实现，保持 external commit 的关键语义： + +- `RMSNorm` / `RMSNormNoWeight` +- RoPE +- Grouped Query Self-Attention +- SwiGLU FFN +- AttnRes operator / residual source aggregation +- `AttnResTransformerBackbone` + +并保持： + +- **full attention**（不使用因果注意力） +- `backbone_type='attnres_full'` +- 输出仅切回 action token 部分，再经过最终 norm + head 得到 velocity-like 输出 + +### 4. Training Objective + +训练目标从当前 DDPM epsilon prediction 改为 external IMF 目标： + +给定真实轨迹 `x` 与随机噪声 `e`： + +1. 采样 `t ~ U(0,1)`、`r ~ U(0,1)`，并排序为 `t >= r` +2. 构造插值状态： + - `z_t = (1 - t) x + t e` +3. 用模型计算： + - `v = f(z_t, t, t, cond)` +4. 对 `g(z, r, t) = f(z, r, t, cond)` 做 JVP，得到： + - `u, du_dt` +5. 构造 compound velocity： + - `V = u + (t - r) * du_dt` +6. 目标为： + - `target = e - x` +7. 用 action 维度上的 MSE 作为最终损失 + +RoboIMI 现有 batch 中的 `action_is_pad` 仍要保留支持；如果存在 padding，只在有效 action 上计算损失。 + +### 5. One-Step Inference + +推理改为 external IMF 的一步采样语义： + +1. 从标准高斯初始化 action trajectory `z_t` +2. 计算 `u = f(z_t, r=0, t=1, cond)` +3. 一步更新： + - `x_hat = z_t - (t-r) * u = z_t - u` +4. 反归一化得到动作序列 + +这意味着： + +- `num_inference_steps` 对 IMF policy 固定为 `1` +- 不再调用 DDIM scheduler 的多步 `step()` +- 在线控制中仍沿用当前 chunk 机制： + - 动作队列为空时触发一次 `predict_action_chunk()` + - 取预测序列中 `[obs_horizon-1 : obs_horizon-1+num_action_steps]` 这一段入队 + +也就是说，**触发模型前向的规则不变，改变的是每次触发后的动作序列生成方式**。 + +## API / Code Structure + +计划中的主要代码边界如下： + +- `roboimi/vla/models/heads/attnres_transformer_components.py` + - IMF AttnRes 基础组件 +- `roboimi/vla/models/heads/imf_transformer1d.py` + - RoboIMI 版本 IMF transformer head + - 对外暴露 `forward(sample, r, t, cond=None)` + - 暴露 `get_optim_groups()` 供 AdamW 分组使用 +- `roboimi/vla/agent_imf.py` + - 复用 `VLAAgent` 的观测处理 / normalization / queue 基础设施 + - 覆盖 IMF 的训练损失与 one-step 预测逻辑 +- Hydra config + - `roboimi/vla/conf/head/imf_transformer1d.yaml` + - `roboimi/vla/conf/agent/resnet_imf_attnres.yaml` + +训练脚本主流程尽量不改；只要求它能 instantiate 新 agent 并继续使用当前 rollout / checkpoint / swanlab 逻辑。 + +## Compatibility Decisions + +## Initial Config Defaults To Preserve + +为避免迁移时语义漂移，首版 IMF 配置默认值明确固定为： + +- `backbone_type: attnres_full` +- `n_head: 1` +- `n_kv_head: 1` +- `n_cond_layers: 0` +- `time_as_cond: true` +- `causal_attn: false` +- `num_inference_steps: 1` + +这些默认值与 external `185ed659` 的 IMF-AttnRes 使用方式保持一致；后续调参可以覆盖，但首版迁移必须先以该语义跑通。 + +### Reuse From RoboIMI + +保留： + +- 三相机数据读取方式 +- ResNet visual backbone +- qpos / action normalization +- 训练循环、优化器、scheduler、SwanLab、headless rollout +- `select_action()` 的在线 chunk 执行方式 + +### Replace With External IMF Semantics + +替换： + +- transformer head 实现 +- diffusion training objective +- inference sampling semantics + +### Intentionally Not Mirrored 1:1 + +不强行与 external repo 一致的部分： + +- external repo 的整体 policy 基类继承体系 +- external repo 的 obs encoder 模块树 +- external repo 的 normalizer / mask generator 框架 + +原因是当前 RoboIMI 已有稳定的数据接口和 rollout 流程，直接嫁接进去更稳。 + +## Testing / Verification Strategy + +迁移完成后至少验证以下内容： + +1. **单元 / 冒烟验证** + - IMF head 前向 shape 正确 + - IMF agent `compute_loss()` 在真实 batch 上可前向、反向 + - IMF agent `predict_action()` 能输出 `(B, pred_horizon, action_dim)` +2. **训练链路验证** + - 使用 GPU 跑一个短训练任务，确认： + - dataloader 正常 + - optimizer / lr scheduler 正常 + - SwanLab 正常记录配置和训练指标 +3. **rollout 验证** + - 训练中周期性 headless rollout 能跑通 + - 环境仍按 EE-style `step()` 接收动作 +4. **最终交付** + - 用用户指定的同类超参数启动正式训练 + +## Risks and Mitigations + +### Risk 1: JVP 在 CUDA 注意力内核上不稳定 + +缓解：沿用 external repo 的策略，在 JVP 路径上切换到 math SDP kernel，必要时 fallback 到 `torch.autograd.functional.jvp`。同时，JVP 的切线构造与 `u, du_dt` 计算流程必须严格对齐 external source，不在本次迁移中自行改写其数学语义。 + +### Risk 2: Optimizer 参数分组遗漏新模块 + +缓解：IMF head 提供 `get_optim_groups()`，并在训练脚本中按“只要 head 提供该接口就使用”的策略统一处理，而不是绑定旧 `head_type`。 + +### Risk 3: 现有 rollout 逻辑假定 DDIM 多步采样 + +缓解：保持 `select_action()` / `predict_action_chunk()` 接口不变，只替换 `predict_action()` 内部实现，确保 eval 代码无需理解 IMF 细节。 + +### Risk 4: 训练命令参数与新 config 不一致 + +缓解：新增独立 agent config，并保留此前训练参数作为显式 CLI override 模板。 + +## Success Criteria + +以下条件全部满足，视为本次迁移成功： + +1. RoboIMI 中新增 IMF-AttnRes policy，可通过 Hydra config 单独启用。 +2. 训练时使用 external IMF 的 loss，而不是当前 DDPM epsilon loss。 +3. 推理时使用 one-step IMF 采样，而不是 DDIM 多步采样。 +4. 三相机图像始终作为条件输入参与模型前向。 +5. 在线 rollout 能在 headless 仿真环境中跑通。 +6. 能按最近一次实验参数模板成功启动训练。 diff --git a/docs/superpowers/specs/2026-04-02-imf-rollout-trajectory-images-design.md b/docs/superpowers/specs/2026-04-02-imf-rollout-trajectory-images-design.md new file mode 100644 index 0000000..3ccf74e --- /dev/null +++ b/docs/superpowers/specs/2026-04-02-imf-rollout-trajectory-images-design.md @@ -0,0 +1,75 @@ +# IMF Rollout Trajectory Images + Short-Horizon Training Design + +## Background +The current RoboIMI IMF training flow can perform rollout validation and log scalar reward metrics to SwanLab, but it does not yet emit the qualitative rollout artifacts now required for analysis. The user wants training-time rollout validation to save front-view trajectory images with the model-generated trajectory drawn in red, upload those images to SwanLab, and then start a new local short-horizon IMF training run. + +## Goals +1. During training-time rollout validation, save one **front-camera** trajectory image per rollout episode. +2. The image must show the rollout EE trajectory in red. +3. Reuse the existing repository trajectory visualization logic as much as practical, especially the existing red capsule-marker trajectory representation. +4. Save 5 rollout images locally for each validation event and upload the same 5 images to SwanLab. +5. Do **not** record rollout videos for this training-time validation flow. +6. Start a new local IMF-AttnRes training run with: + - `agent.head.n_emb=384` + - `agent.head.n_layer=12` + - `agent.pred_horizon=8` + - `agent.num_action_steps=4` + - `train.max_steps=50000` + - `train.rollout_num_episodes=5` + - `train.use_swanlab=true` + +## Non-Goals +- No IMF architecture or loss-function change. +- No dataset schema change. +- No rollout video generation for the new training flow. +- No interactive viewer requirement. + +## Existing Relevant Code +- `roboimi/demos/vla_scripts/eval_vla.py` + - already supports rollout summaries, optional trajectory export, and optional video export. +- `roboimi/utils/raw_action_trajectory_viewer.py` + - already contains the red trajectory capsule-marker construction logic. +- `roboimi/demos/vla_scripts/train_vla.py` + - already performs periodic rollout validation and scalar SwanLab logging. +- `roboimi/vla/agent.py` + - already implements “predict pred_horizon, execute first num_action_steps” queue semantics. + +## Design Decisions + +### 1. Artifact contract +Each rollout episode will emit one distinct PNG file under the eval artifact directory. The file naming/path contract must be per-episode, not shared, so a 5-episode validation event yields 5 stable image paths without overwriting. + +### 2. Trajectory definition +The red trajectory corresponds to the **actually executed model action sequence** over the rollout loop: the raw EE actions returned and consumed step-by-step by the policy loop. For the requested short-horizon run, this means the visualization reflects repeated execution of the first 4 actions from each predicted 8-action chunk, not every discarded future prediction from replanning. + +### 3. Camera choice +The training-time image export path is explicitly pinned to the repo’s concrete `front` camera key. It must not silently use `camera_names[0]` if that is not `front`. + +### 4. Rendering path +`eval_vla.py` will add a lightweight headless image-export path that: +- renders the `front` camera frame, +- overlays the trajectory using the existing red trajectory representation, +- saves a static PNG per episode. + +The implementation may reuse the existing marker-construction logic directly and add a minimal helper for final image composition/export. + +### 5. Training-time behavior +`train_vla.py` rollout validation must explicitly: +- request/save trajectory images, +- keep `record_video=false`, +- return the 5 per-episode image paths in the rollout summary payload, +- upload those 5 images to SwanLab, +- keep image-upload failures non-fatal. + +## Expected User-Visible Outcome +For each scheduled validation event in the new training run: +- 5 rollout episodes execute, +- 5 front-view PNG trajectory images are saved locally, +- the same 5 images are uploaded to SwanLab, +- scalar reward metrics continue to be logged, +- no rollout videos are generated. + +## Risks and Mitigations +- **Headless rendering conflicts from desktop env vars**: force headless eval onto EGL when `headless=true`. +- **Image overwrite risk**: use explicit per-episode artifact paths. +- **SwanLab media API mismatch**: isolate media logging in a small best-effort helper. diff --git a/docs/superpowers/specs/2026-04-05-lewm-vit-backbone-design.md b/docs/superpowers/specs/2026-04-05-lewm-vit-backbone-design.md new file mode 100644 index 0000000..96511a9 --- /dev/null +++ b/docs/superpowers/specs/2026-04-05-lewm-vit-backbone-design.md @@ -0,0 +1,138 @@ +# LEWM ViT Backbone Replacement Design + +## Goal +将当前 roboimi VLA policy 中的 ResNet 视觉编码器替换为来自 LEWM checkpoint 的冻结 ViT 视觉编码器（encoder + projector），仅使用最终 CLS token 的 192 维 embedding 作为视觉特征。 + +## User constraints +- 使用 `/home/droid/下载/lewm_sim_transfer_checkpoint_usage.md` 中确认的训练好 checkpoint +- 只使用视觉编码部分：`encoder + projector` +- 权重冻结 +- 维持“视觉特征 + state 拼接，再送入 diffusion transformer”这一总体处理方式 +- 输入使用三视角：`[r_vis, top, front]` +- 在 5880 机器上启动两个训练：`embed=384/layer=12` 和 `embed=256/layer=12` +- `pred_horizon=16` +- `num_action_steps=8` +- 每个训练 `50k` steps +- rollout 验证每次用 `10` 个 episodes，不是之前的 `5` + +## Trusted existing facts +1. LEWM checkpoint 路径： + - `/home/droid/le-wm/lewm-sim-transfer/pa1w85md8jop6bvol8oxp/checkpoints/epoch=99-step=47800.ckpt` +2. 需要加载的 state_dict 前缀： + - `model.encoder.*` + - `model.projector.*` +3. LEWM ViT 配置： + - encoder scale: `tiny` + - hidden size: `192` + - layers: `12` + - attention heads: `3` + - patch size: `14` + - projector: `MLP(192 -> 2048 -> 192)` with `BatchNorm1d + GELU` +4. LEWM 训练时三视角先拼成单图，再送入单个 ViT encoder；输出整体视觉 embedding 是 **192 维**。 + +## Key design decision +### Chosen design: fuse 3 cameras into one LEWM-style image, output one 192-d visual vector per timestep +不是把 LEWM ViT 当成“每相机一个 192-d encoder”，而是按 LEWM 原训练方式： +- 输入三视角图像字典 `{r_vis, top, front}` +- 按固定顺序拼成一张 fused image +- 走单个 frozen ViT + projector +- 得到一个 **192 维总视觉特征** + +### Why this is the right replacement +当前 ResNet backbone 对外给到 policy head 的**总视觉特征维度**是： +- 每相机 `64` +- 三相机总计 `192` + +而 LEWM checkpoint 输出的 CLS/projector embedding 也是： +- 总计 `192` + +因此，最自然的“直接平替当前 ResNet 视觉编码器”的方式是： +- 用 LEWM backbone 直接产出一个 192-d 总视觉向量 +- 后续和 state `16-d` 拼接后，依旧得到 `208-d` 条件向量 +- 不改 diffusion head 的总体接口和语义 + +## Interface compatibility plan +现有 `VLAAgent` 假设 backbone 暴露： +- `camera_names` +- `num_cameras` +- `output_dim`（语义上是“每相机特征维度”） +- `forward(images_dict) -> (B, T, total_visual_dim)` + +为了最小改动兼容现有 agent： +- 新 LEWM backbone 的 `forward()` 返回 `(B, T, 192)` +- `camera_names = ('r_vis', 'top', 'front')` +- `num_cameras = 3` +- `output_dim = 64` + +这样 `VLAAgent` 内部仍会计算： +- `per_step_cond_dim = output_dim * num_cams + obs_dim = 64*3 + 16 = 208` +与实际 `forward()` 输出的 `192 + 16 = 208` 保持一致。 + +> 也就是说：`output_dim` 在这个 backbone 里保留为“与旧 ResNet 总特征等价的单相机占位维度”，而不是“真实 projector 输出维度”。这是一个兼容性 shim，用来避免改 agent 主逻辑。 + +## Image preprocessing design +当前 roboimi dataset 已经把每个相机图像读成： +- `(C, 224, 224)` +- 值域 `[0, 1]` + +新 LEWM backbone 将： +1. 按顺序取 `r_vis`, `top`, `front` +2. 在宽度方向拼接，得到 fused image： + - `(C, 224, 672)` +3. 使用 LEWM 一致的 ImageNet normalize： + - mean `[0.485, 0.456, 0.406]` + - std `[0.229, 0.224, 0.225]` +4. 调用 `ViTModel(..., interpolate_pos_encoding=True)` +5. 取 `last_hidden_state[:, 0]` +6. 送入 frozen projector，得到 `(B*T, 192)` + +## Files to create / modify +### New files +- `roboimi/vla/models/backbones/lewm_vit_backbone.py` +- `roboimi/vla/conf/backbone/lewm_vit_diffusion.yaml` +- `roboimi/vla/conf/agent/lewm_imf_attnres.yaml` +- `tests/test_lewm_vit_backbone.py` + +### Modified files +- `roboimi/vla/models/backbones/__init__`（如果需要导出） +- `tests/test_imf_vla_agent.py`（增加新 backbone 集成用例） +- `roboimi/demos/vla_scripts/train_vla.py`（如需仅调整 rollout 默认/日志；如果命令覆盖足够，则尽量不改主逻辑） +- 训练/实验 suite 文档（新增本次 LEWM ViT 训练记录） + +## Testing plan +1. **Unit test: load + forward** + - 用 synthetic checkpoint 验证新 backbone 能正确加载 `model.encoder.*` 与 `model.projector.*` + - 输入 3 相机 `(B,T,C,224,224)` + - 输出 `(B,T,192)` +2. **Agent integration test** + - backbone.output_dim=64, num_cameras=3 + - agent `_build_cond()` 输出最后维度为 `208` +3. **Remote smoke test on 5880** + - 使用真实 checkpoint + - `max_steps=2` + - 两个实验各自 smoke 一次 +4. **Full run** + - GPU0: `embed=384, layer=12` + - GPU1: `embed=256, layer=12` + - `rollout_num_episodes=10` + +## Training launch contract +- host: `100.73.14.65` +- code dir: `/home/droid/roboimi_suite_20260404` +- python: `/home/droid/miniforge3/envs/roboimi/bin/python` +- dataset: `/home/droid/sim_dataset/sim_transfer` +- cameras: `[r_vis, top, front]` +- agent: new `lewm_imf_attnres` +- max_steps: `50000` +- rollout every `5` epochs +- rollout episodes: `10` + +## Risks +1. LEWM 训练时的 fused image 预处理如果方向实现错了（224x672 vs 672x224），会导致分布偏移。 +2. 当前 roboimi env 需确保安装 `transformers`；从 `environment.yml` 看本地已有该依赖，但远端训练环境要 smoke 确认。 +3. 因为这是 frozen ViT + projector，若 projector BN 仍保持 train 模式，统计量会漂移，所以必须整体 `eval()` 并冻结。 + +## Recommended first implementation path +- 先实现一个独立 `LEWMViTBackbone` 类，不改现有 `ResNetDiffusionBackbone` 主逻辑。 +- 再通过新的 hydra backbone/agent 配置接入。 +- 优先做到“最少侵入 + smoke 可跑 + 远端可训”。 diff --git a/docs/superpowers/specs/2026-04-05-phase2-full-attnres-vision-design.md b/docs/superpowers/specs/2026-04-05-phase2-full-attnres-vision-design.md new file mode 100644 index 0000000..b1c2f0c --- /dev/null +++ b/docs/superpowers/specs/2026-04-05-phase2-full-attnres-vision-design.md @@ -0,0 +1,81 @@ +# Phase-2 Full-AttnRes Vision Design + +## Goal +在当前 roboimi IMF policy 中，把视觉 backbone 里原先由 ResNet BasicBlock/Bottleneck 提供的残差单元全部替换为 AttnRes 风格单元，同时尽量保持现有 agent / cond / rollout / 训练脚本接口不变。 + +## User requirement interpretation +这里按最严格解释执行： +- 不是“在 ResNet 后面再加一个 AttnRes 模块” +- 也不是“只在某几个 stage 加 AttnRes 混合” +- 而是：视觉主干网络中原本依赖 ResNet residual block 的地方，统一改成 AttnRes residual operator 驱动的 block +- 最终仍然输出与现有 `ResNetDiffusionBackbone` 相同的每相机特征接口，以便复用 `SpatialSoftmax -> Linear -> ReLU`、多相机拼接、state concat、IMF head 条件输入 + +## Recommended design +### Option A (recommended) +保留 ResNet 的宏观 stage/stem 结构与通道/步幅规划，但把每个 stage 内的 BasicBlock/Bottleneck 替换为新的 `AttnResImageBlock2D`： +- 输入仍是 `(B, C, H, W)` feature map +- block 内先把空间维 flatten 成 token 序列 `(B, H*W, C)` +- 用二维位置编码 / 可学习位置偏置 + AttnRes self-attention + AttnRes FFN 完成 block 变换 +- 再 reshape 回 `(B, C, H, W)` +- stage 间下采样仍由显式 stride/downsample path 完成 + +优点： +- 最接近“ResNet 中所有残差都由 AttnRes 代替”的要求 +- 保留现有视觉输出接口和 cond_dim，不用改 agent/head/data pipeline +- 仍可沿用现有多相机编码器框架 + +缺点： +- 需要新写 2D 版 AttnRes image block，而不是直接复用 1D IMF head block + +### Option B +完全移除 ResNet stage，换成 patchify + ViT/AttnRes 图像 transformer，再接 SpatialSoftmax/MLP。 + +优点：实现概念更统一。 +缺点：已经不算“把 ResNet 中残差替换掉”，而是直接换 backbone，和用户要求不完全一致。 + +### Option C +保留现有 ResNet block，只在 block 外层加 AttnRes mixing。 + +不推荐，因为不满足“所有残差均由 AttnRes 替代”。 + +## Concrete architecture choice +采用 Option A： +1. 保留 stem（conv/bn-or-gn/relu/maxpool）与 stage 边界 +2. 新增 `AttnResImageBlock2D` +3. 新增 `AttnResResNetLikeBackbone2D`，负责堆叠 stage/block +4. 在 `ResNetDiffusionBackbone` 中增加可选 backbone mode，例如： + - `vision_backbone_mode: resnet` + - `vision_backbone_mode: attnres_resnet` +5. `resnet_imf_attnres` agent 配置新增一个 Phase-2 变体，默认打开 `attnres_resnet` +6. 仍保持： + - 每相机输出 `64` + - 多相机总视觉输出 `3 * 64` + - 与 state 拼接后 `cond_dim = 208` + +## Files likely to change +- `roboimi/vla/models/backbones/resnet_diffusion.py` +- `roboimi/vla/conf/backbone/resnet_diffusion.yaml` +- `roboimi/vla/conf/agent/resnet_imf_attnres.yaml` +- new: `roboimi/vla/models/backbones/attnres_resnet2d.py` +- tests: + - new: `tests/test_attnres_resnet2d_backbone.py` + - update/add wiring test for agent cond dims + +## Test plan +1. New backbone instantiates and forwards `(B,T,C,H,W)` multi-camera input +2. Output shape unchanged vs current backbone +3. `output_dim == 64` +4. 3-camera cond path still yields `208` +5. Phase-2 config instantiates full IMF agent successfully +6. One short CPU smoke forward for `compute_loss` + +## Phase-2 experiment plan +固定使用 Phase-1 最优组合： +- `pred_horizon=16` +- `num_action_steps=8` + +比较： +1. baseline: current IMF head-only AttnRes + original ResNet vision backbone +2. phase2: IMF head AttnRes + full AttnRes-replaced vision backbone + +训练超参保持与 Phase-1 最优设置一致，先跑一组 50k step 对比。 diff --git a/docs/superpowers/specs/2026-04-06-resnet-multitoken-imf-design.md b/docs/superpowers/specs/2026-04-06-resnet-multitoken-imf-design.md new file mode 100644 index 0000000..3bcf5d5 --- /dev/null +++ b/docs/superpowers/specs/2026-04-06-resnet-multitoken-imf-design.md @@ -0,0 +1,32 @@ +# ResNet Multitoken IMF Design + +**Status:** user-specified architecture, treated as approved on 2026-04-06. + +## Goal +Keep a standard ResNet-18 visual trunk (no AttnRes in vision), but change IMF conditioning from one concatenated multiview token per obs step into three camera-specific condition tokens per obs step. + +## Approved architecture +- Vision trunk: standard `resnet18` residual network +- Cameras: `front`, `top`, `r_vis` +- Each camera uses its **own** ResNet-18 weights (`use_separate_rgb_encoder_per_camera=true`) +- Each camera produces one visual token +- For each obs step and each camera: + 1. take that camera visual token + 2. concatenate robot state + 3. project to one condition token +- IMF input should receive **3 condition tokens per obs step**, not one concatenated token +- With `obs_horizon=2`, IMF cond sequence length becomes `2 * 3 = 6` +- IMF head remains on the existing IMF/AttnRes implementation path +- Vision trunk remains standard ResNet; **no AttnRes vision replacement** + +## Design choices +- Extend `ResNetDiffusionBackbone` with an opt-in mode that returns per-camera tokens shaped `(B, T, num_cams, D)` instead of concatenating camera features into `(B, T, num_cams * D)`. +- Teach `VLAAgent` to detect multi-token visual features, broadcast state per camera token, apply the existing condition projector on each token, then flatten `(T, num_cams)` into one cond sequence for the IMF head. +- Keep `per_step_cond_dim` as the width of a single condition token, and add explicit token-count metadata so transformer heads get the correct cond-sequence length. +- For the new experiments, set the condition-token width equal to `n_emb` via `cond_projector.output_dim=${agent.head.n_emb}`. + +## Files expected to change +- `roboimi/vla/models/backbones/resnet_diffusion.py` +- `roboimi/vla/agent.py` +- new Hydra agent config for the multitoken ResNet IMF variant +- focused tests in `tests/test_imf_vla_agent.py` and/or `tests/test_resnet_transformer_agent_wiring.py` diff --git a/docs/superpowers/specs/2026-04-06-siglip2-multiview-vla-design.md b/docs/superpowers/specs/2026-04-06-siglip2-multiview-vla-design.md new file mode 100644 index 0000000..e2c27da --- /dev/null +++ b/docs/superpowers/specs/2026-04-06-siglip2-multiview-vla-design.md @@ -0,0 +1,41 @@ +# SigLIP2 Multiview VLA Design + +**Status:** user-specified architecture, treated as approved on 2026-04-06 + +## Goal +Replace the current vision encoder for the IMF/AttnRes diffusion policy with a frozen SigLIP2 image encoder while preserving the downstream action-diffusion stack and rollout behavior. + +## Approved architecture +- Backbone model: `google/siglip2-base-patch16-256` +- Camera inputs: three views, encoded **independently** with a **shared** SigLIP2 vision encoder +- Input size: + - dataset images stay at native `256x256` (no dataset-side resize) + - eval/rollout images resize to `256x256` before SigLIP2 because env renders are larger +- Per-view feature: use the global pooled image feature (`pooler_output`, 768-d) +- Per-view projection experiments: + 1. `768 -> 96` + 2. `768 -> 192` +- Conditioning pipeline: + 1. concatenate 3 projected camera vectors + 2. concatenate robot state + 3. project concatenated condition to `384` + 4. feed that `384`-d per-step condition into the existing IMF/AttnRes diffusion head +- Training/run defaults for requested experiments: + - `n_emb=384` + - `n_layer=12` + - `pred_horizon=16` + - `num_action_steps=8` + - rollout count for validation: keep current requested behavior on this branch unless explicitly overridden later + +## Design decisions +- The condition projector lives in `VLAAgent._build_cond()` so the backbone owns only visual features, while the agent owns the final conditioning contract expected by the diffusion head. +- The SigLIP2 backbone is frozen by default; only the per-view projectors and downstream policy layers train. +- The backbone exposes `dataset_image_resize_shape=None` and `eval_image_resize_shape=(256, 256)` so existing train/eval plumbing can reuse the raw-256 path already added in this branch. +- One shared vision encoder is used across cameras to keep memory and download size reasonable and to match the user's request for per-view independent encoding rather than a fused multiview image. + +## Files expected to change +- `roboimi/vla/models/backbones/` for the new SigLIP2 backbone +- `roboimi/vla/agent.py` for optional post-concat condition projection +- Hydra configs under `roboimi/vla/conf/{agent,backbone,modules}` +- tests for backbone wiring and agent conditioning dims +- remote launch commands/scripts only as needed for training diff --git a/docs/superpowers/specs/2026-04-23-sim-air-insert-ring-bar-design.md b/docs/superpowers/specs/2026-04-23-sim-air-insert-ring-bar-design.md new file mode 100644 index 0000000..feb54b6 --- /dev/null +++ b/docs/superpowers/specs/2026-04-23-sim-air-insert-ring-bar-design.md @@ -0,0 +1,316 @@ +# sim_air_insert_ring_bar Design + +## Summary + +Add a new independent MuJoCo simulation task named `sim_air_insert_ring_bar` that keeps the existing dual-Diana tabletop setup but replaces the single transfer box with two randomized objects: + +- a square ring block grasped by the left arm +- a square bar block grasped by the right arm + +The task is to pick both objects off the table and complete an in-air insertion where the bar truly passes through the ring aperture. The existing `sim_transfer` task must remain unchanged. + +## Goals + +- Reuse the current dual-Diana EE-control simulation stack +- Keep the same table/base robot arrangement as the existing transfer task +- Add an independent task entrypoint and scene definition +- Randomize planar placement of both objects within left/right task-specific regions +- Implement reward staging for contact, lift, and successful in-air insertion +- Add a scripted policy that performs pick, lift, align, and in-air insertion +- Preserve compatibility with existing environment creation, evaluation, and rollout patterns + +## Non-Goals + +- No random yaw in the first version +- No visual servoing or closed-loop insertion controller +- No general multi-task environment framework refactor +- No guarantee that the VLA training stack is immediately tuned for this new task +- No replacement or behavior change for `sim_transfer` + +## Task Name + +Use a new task name: + +- `sim_air_insert_ring_bar` + +This task should be exposed alongside `sim_transfer`, not as a replacement. + +## Scene Geometry + +### Shared Base Scene + +Keep the dual Diana robot, the table, and the existing camera layout conceptually unchanged. + +### Ring Block + +Represent the square ring as a rigid free body composed from simple MuJoCo box geoms rather than an external mesh. + +Dimensions: + +- outer side length: 68 mm +- inner aperture side length: 32 mm +- thickness: 18 mm +- ring wall width: 18 mm + +The ring should behave as a single object body with a single free joint. + +### Bar Block + +Represent the bar as a rigid free body with a single box geom. + +Dimensions: + +- length: 90 mm +- cross-section: 18 mm x 18 mm + +The bar should also be a single free-joint body. + +## Initial Placement / Reset + +The first version uses position-only randomization with fixed orientation. Reset sampling stays **caller-driven**, matching the existing `sim_transfer` usage pattern in rollout/eval code: a helper samples task state, then callers pass that state into `env.reset(...)`. + +Use an explicit sampled task-state structure with named fields: + +- `ring_pos`: 3D position +- `ring_quat`: fixed 4D quaternion for version 1 +- `bar_pos`: 3D position +- `bar_quat`: fixed 4D quaternion for version 1 + +Behavior: + +- ring block: randomized only in a left-side planar sampling region +- bar block: randomized only in a right-side planar sampling region +- both objects start flat on the table +- both objects use fixed orientation at reset +- no random yaw, tilt, or flip in this version + +The sampling regions should be chosen conservatively so that: + +- the left arm can comfortably reach and grasp the ring +- the right arm can comfortably reach and grasp the bar +- scripted open-loop pick trajectories remain feasible + +## Control / Action Interface + +Reuse the current 16D EE-space action convention already used by the dual-Diana position-control environment: + +- left arm EE pose: 7D (`xyz + quat`) +- right arm EE pose: 7D (`xyz + quat`) +- left gripper command: 1D +- right gripper command: 1D + +The new task should continue using EE targets transformed through the existing IK-based control path. + +## Environment Structure + +Implement this as a new task-specific environment path while reusing the existing dual-Diana simulation base where possible. + +Expected responsibilities: + +- scene instantiation for the ring+bar setup +- task reset for randomized object placement +- environment-state accessors for both objects +- reward computation +- in-air insertion success detection + +The environment factory must dispatch by task name and leave the `sim_transfer` branch unchanged. + +## Observation / Environment State + +The task should retain the current observation structure style used by the dual-Diana environment: + +- `qpos` +- multi-camera images + +For task state access, the environment should expose a stable `env_state` vector with this exact order: + +- `ring_pos[0:3]` +- `ring_quat[3:7]` +- `bar_pos[7:10]` +- `bar_quat[10:14]` + +This 14D state should be sufficient for scripted-policy debugging and future rollout analysis, while reset itself remains caller-driven via the named task-state helper structure above. + +## Reward Design + +Use staged rewards in the same spirit as the current task, returning the highest achieved stage rather than accumulating one-time sparse bonuses per event. + +Maximum reward: + +- `max_reward = 5` + +Reward stages: + +1. left gripper touches the ring block +2. right gripper touches the bar block +3. ring block is lifted off the table +4. bar block is lifted off the table +5. while both objects are off the table, the bar truly passes through the ring aperture + +Notes: + +- contact rewards are intended as grasp-progress stages +- lift rewards require the object to be off the table, not merely touched +- final success reward only applies when both objects are airborne + +## Success Detection + +Success must **not** be based on a centerline-only check. + +A centerline-only test is insufficient because: + +- the bar has thickness, so a centerline can pass through while the body cannot +- a square bar with imperfect orientation can have its centerline inside the aperture while its corners still collide with the ring + +### Required Success Semantics + +A successful insertion requires all of the following: + +1. the ring is off the table +2. the bar is off the table +3. the bar has actually crossed through the ring thickness direction +4. the bar’s finite square cross-section fits through the square aperture during that crossing + +### Recommended Detection Approach + +Use a task-level geometric check in Python rather than relying on contact alone. + +Implementation intent: + +- transform the bar geometry into the ring’s local frame +- reason about the bar as a finite oriented box (not a line) +- verify that the bar has crossed the ring thickness direction +- verify that the portion of the bar passing the aperture fits within the inner square opening, accounting for the bar’s cross-section and orientation + +This geometric check is the primary success test. + +### Role of Contacts + +Contacts may still be used for: + +- grasp-stage rewards +- debugging / diagnostics + +But contact alone should **not** be the sole criterion for insertion success, since: + +- a true clean insertion may have limited aperture-wall contact +- persistent contact can also happen while the bar is jammed and not actually inserted + +## Scripted Policy + +Add a new task-specific scripted policy for `sim_air_insert_ring_bar`. + +### Policy Intent + +The first version prioritizes a conservative, reliable open-loop demonstration rather than an optimized trajectory. + +### Action Phases + +The scripted policy should follow these phases: + +1. move both arms to safe initial / waiting poses with grippers open +2. move left arm above the ring and right arm above the bar +3. descend and grasp the assigned objects +4. lift both objects clear of the table +5. move both objects to an airborne meeting region above the table +6. hold the ring stably while aligning the bar with the aperture +7. push the bar along the intended insertion direction until the geometric success condition is met + +### Grasp Assignment + +- left arm: ring only +- right arm: bar only + +### Motion Style + +Keep the current repository style: + +- waypoint-based trajectory definition +- open-loop interpolation between waypoints +- fixed grasp orientation in the first version + +No adaptive replanning is required for the first version. + +## Files / Integration Scope + +The implementation is expected to add task-specific files rather than broadly refactoring the codebase. + +Likely additions / changes: + +- a new MuJoCo scene XML for the ring+bar task +- one or more XML fragments defining the two new objects +- a new task-specific dual-Diana environment file +- robot asset wiring for the new scene XML +- reset sampling helpers for the new task +- task registration in constants / environment factory paths +- a new scripted policy file +- focused tests for task creation, reset, rewards, success detection, and scripted policy shape/smoke behavior + +## Testing Requirements + +At minimum, add regression coverage for: + +### Environment Creation + +- the new task can be created via the task factory +- the existing `sim_transfer` task remains unchanged + +### Reset / Sampling + +- ring reset positions are inside the left sampling region +- bar reset positions are inside the right sampling region +- reset orientation is fixed as intended + +### Environment State + +- environment-state access returns both object poses in the expected structure + +### Success Detection + +Must include both positive and negative cases. + +Positive case: + +- a configuration where the finite bar truly passes through the ring aperture is detected as success + +Negative cases: + +- centerline-inside but finite body would clip the aperture +- not enough depth / not actually crossing the ring thickness direction +- one or both objects still on the table + +### Reward Logic + +- left contact stage +- right contact stage +- ring lift stage +- bar lift stage +- final success stage with `max_reward = 5` + +### Scripted Policy + +At minimum: + +- policy emits valid 16D actions +- trajectory generation does not error +- rollout smoke path can step through the new environment + +## Risks / Constraints + +- MuJoCo contact naming must remain stable enough for stage rewards +- geometric insertion checks must be strict enough to avoid false positives but not so brittle that numerically valid insertions are missed +- scripted open-loop insertion may require conservative alignment and lift heights to keep the first version reliable + +## Acceptance Criteria + +The feature is complete when all of the following are true: + +- `sim_air_insert_ring_bar` is creatable as an independent task +- the scene contains the dual Diana, table, ring block, and bar block +- reset randomizes ring and bar positions in left/right planar regions with fixed orientation +- the environment exposes task state for both objects +- staged rewards progress to `max_reward = 5` +- final success is based on finite-geometry insertion semantics, not a centerline-only shortcut +- a new scripted policy can execute the intended pick-lift-align-insert behavior in the new environment +- a canonical nominal smoke path (unit-level or deterministic integration-level) exists for the new scripted-policy interface so success is not judged purely by interpretation +- existing `sim_transfer` behavior is preserved diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..7f1c879 --- /dev/null +++ b/environment.yml @@ -0,0 +1,458 @@ +name: roboimi +channels: + - conda-forge +dependencies: + - _libgcc_mutex=0.1 + - _openmp_mutex=4.5 + - _python_abi3_support=1.0 + - aiohappyeyeballs=2.6.1 + - aiohttp=3.13.3 + - aiosignal=1.4.0 + - alsa-lib=1.2.9 + - anyio=4.12.1 + - aom=3.5.0 + - async-timeout=5.0.1 + - attr=2.5.1 + - attrs=25.4.0 + - aws-c-auth=0.7.22 + - aws-c-cal=0.6.15 + - aws-c-common=0.9.23 + - aws-c-compression=0.2.18 + - aws-c-event-stream=0.4.2 + - aws-c-http=0.8.2 + - aws-c-io=0.14.9 + - aws-c-mqtt=0.10.4 + - aws-c-s3=0.5.10 + - aws-c-sdkutils=0.1.16 + - aws-checksums=0.1.18 + - aws-crt-cpp=0.26.12 + - aws-sdk-cpp=1.11.329 + - box2d-py=2.3.8 + - brotli=1.1.0 + - brotli-bin=1.1.0 + - brotli-python=1.1.0 + - bzip2=1.0.8 + - c-ares=1.34.6 + - ca-certificates=2026.1.4 + - cairo=1.16.0 + - certifi=2026.1.4 + - cffi=1.17.1 + - charset-normalizer=3.4.4 + - click=8.3.1 + - cloudpickle=3.0.0 + - contourpy=1.3.0 + - cpython=3.10.19 + - cuda-cudart=12.6.68 + - cuda-cudart_linux-64=12.6.68 + - cuda-nvrtc=12.6.68 + - cuda-nvtx=12.6.68 + - cuda-version=12.6 + - cudnn=8.9.7.29 + - cycler=0.12.1 + - datasets=4.0.0 + - dav1d=1.2.1 + - dbus=1.13.6 + - dill=0.3.8 + - eigen=3.4.0 + - exceptiongroup=1.3.1 + - expat=2.6.3 + - farama-notifications=0.0.4 + - filelock=3.15.4 + - fluidsynth=2.3.3 + - font-ttf-dejavu-sans-mono=2.37 + - font-ttf-inconsolata=3.000 + - font-ttf-source-code-pro=2.038 + - font-ttf-ubuntu=0.83 + - fontconfig=2.14.2 + - fonts-conda-ecosystem=1 + - fonts-conda-forge=1 + - fonttools=4.53.1 + - freetype=2.12.1 + - frozenlist=1.7.0 + - fsspec=2024.6.1 + - gettext=0.22.5 + - gettext-tools=0.22.5 + - gflags=2.2.2 + - git-lfs=3.7.1 + - glog=0.7.1 + - gmp=6.3.0 + - gmpy2=2.1.5 + - graphite2=1.3.13 + - gym=0.26.1 + - gym-box2d=0.26.1 + - gym-notices=0.0.8 + - gymnasium=0.29.1 + - h11=0.16.0 + - h2=4.3.0 + - harfbuzz=7.3.0 + - hf-xet=1.2.1 + - hpack=4.1.0 + - httpcore=1.0.9 + - httpx=0.28.1 + - huggingface_hub=1.3.5 + - hyperframe=6.1.0 + - icu=72.1 + - idna=3.11 + - jack=1.9.22 + - jax-jumpy=1.0.0 + - jinja2=3.1.4 + - jpeg=9e + - keyutils=1.6.3 + - kiwisolver=1.4.9 + - krb5=1.21.3 + - lame=3.100 + - lcms2=2.15 + - ld_impl_linux-64=2.40 + - lerc=4.0.0 + - libabseil=20240116.2 + - libarrow=16.1.0 + - libarrow-acero=16.1.0 + - libarrow-dataset=16.1.0 + - libarrow-substrait=16.1.0 + - libasprintf=0.22.5 + - libasprintf-devel=0.22.5 + - libavif=0.11.1 + - libblas=3.9.0 + - libbrotlicommon=1.1.0 + - libbrotlidec=1.1.0 + - libbrotlienc=1.1.0 + - libcap=2.69 + - libcblas=3.9.0 + - libcrc32c=1.1.2 + - libcublas=12.6.1.4 + - libcufft=11.2.6.59 + - libcurand=10.3.7.68 + - libcurl=8.12.1 + - libcusolver=11.6.4.69 + - libcusparse=12.5.3.3 + - libdb=6.2.32 + - libdeflate=1.17 + - libedit=3.1.20250104 + - libev=4.33 + - libevent=2.1.12 + - libexpat=2.6.3 + - libffi=3.4.2 + - libflac=1.4.3 + - libgcc=14.1.0 + - libgcc-ng=14.1.0 + - libgcrypt=1.11.0 + - libgettextpo=0.22.5 + - libgettextpo-devel=0.22.5 + - libgfortran=14.1.0 + - libgfortran-ng=14.1.0 + - libgfortran5=14.1.0 + - libglib=2.80.3 + - libgoogle-cloud=2.25.0 + - libgoogle-cloud-storage=2.25.0 + - libgpg-error=1.50 + - libgrpc=1.62.2 + - libhwloc=2.9.3 + - libiconv=1.17 + - libjpeg-turbo=2.1.4 + - liblapack=3.9.0 + - libmad=0.15.1b + - libmagma=2.8.0 + - libmagma_sparse=2.8.0 + - libnghttp2=1.67.0 + - libnsl=2.0.1 + - libnvjitlink=12.6.68 + - libogg=1.3.5 + - libopenblas=0.3.27 + - libopus=1.3.1 + - libparquet=16.1.0 + - libpng=1.6.43 + - libprotobuf=4.25.3 + - libre2-11=2023.09.01 + - libsndfile=1.2.2 + - libsqlite=3.46.0 + - libssh2=1.11.1 + - libstdcxx=14.1.0 + - libstdcxx-ng=14.1.0 + - libsystemd0=256.5 + - libthrift=0.19.0 + - libtiff=4.5.0 + - libtorch=2.4.0 + - libutf8proc=2.8.0 + - libuuid=2.38.1 + - libuv=1.48.0 + - libvorbis=1.3.7 + - libwebp-base=1.4.0 + - libxcb=1.13 + - libxcrypt=4.4.36 + - libxml2=2.11.5 + - libzlib=1.3.1 + - llvm-openmp=18.1.8 + - lz4-c=1.9.4 + - markupsafe=2.1.5 + - matplotlib-base=3.9.2 + - mkl=2023.2.0 + - mpc=1.3.1 + - mpfr=4.2.1 + - mpg123=1.31.3 + - mpmath=1.3.0 + - multidict=6.7.0 + - multiprocess=0.70.16 + - munkres=1.1.4 + - nccl=2.22.3.1 + - ncurses=6.5 + - networkx=3.3 + - numpy=1.26.4 + - openjpeg=2.5.0 + - openssl=3.6.1 + - opusfile=0.12 + - orc=2.0.1 + - orocos-kdl=1.5.1 + - packaging=24.1 + - pandas=2.2.2 + - pcre2=10.44 + - pillow=9.4.0 + - pip=24.2 + - pixman=0.43.2 + - portaudio=19.6.0 + - portmidi=2.0.4 + - propcache=0.3.1 + - pthread-stubs=0.4 + - pulseaudio-client=16.1 + - pyarrow=16.1.0 + - pyarrow-core=16.1.0 + - pybind11=2.13.5 + - pybind11-global=2.13.5 + - pycparser=2.22 + - pygame=2.1.3 + - pyparsing=3.1.4 + - pysocks=1.7.1 + - python=3.10.14 + - python-dateutil=2.9.0 + - python-gil=3.10.19 + - python-orocos-kdl=1.5.1 + - python-tzdata=2024.1 + - python-xxhash=3.6.0 + - python_abi=3.10 + - pytorch=2.4.0 + - hydra-core=1.3.2 + - omegaconf=2.3.0 + - einops=0.8.2 + - diffusers=0.36.0 + - torchvision=0.19.0 + - pytz=2024.1 + - pyyaml=6.0.3 + - qhull=2020.2 + - re2=2023.09.01 + - readline=8.2 + - regex=2026.1.15 + - requests=2.32.5 + - s2n=1.4.16 + - safetensors=0.7.0 + - sdl2=2.26.5 + - sdl2_image=2.6.3 + - sdl2_mixer=2.6.3 + - sdl2_ttf=2.20.2 + - setuptools=72.2.0 + - shellingham=1.5.4 + - six=1.16.0 + - sleef=3.6.1 + - snappy=1.2.2 + - sniffio=1.3.1 + - stable-baselines3=2.3.2 + - sympy=1.13.2 + - tbb=2021.11.0 + - tk=8.6.13 + - tokenizers=0.22.2 + - tqdm=4.67.2 + - transformers=5.0.0 + - typer-slim=0.21.1 + - typing-extensions=4.12.2 + - typing_extensions=4.12.2 + - tzdata=2024a + - unicodedata2=15.1.0 + - urllib3=2.5.0 + - wheel=0.44.0 + - xorg-kbproto=1.0.7 + - xorg-libice=1.1.1 + - xorg-libsm=1.2.4 + - xorg-libx11=1.8.4 + - xorg-libxau=1.0.11 + - xorg-libxdmcp=1.1.3 + - xorg-libxext=1.3.4 + - xorg-libxrender=0.9.10 + - xorg-renderproto=0.11.1 + - xorg-xextproto=7.3.0 + - xorg-xproto=7.0.31 + - xxhash=0.8.3 + - xz=5.2.6 + - yaml=0.2.5 + - yarl=1.22.0 + - zlib=1.3.1 + - zstandard=0.23.0 + - zstd=1.5.6 + - pip: + - GitPython==3.1.46 + - Jinja2==3.1.6 + - MarkupSafe==3.0.3 + - PyOpenGL==3.1.7 + - PyYAML==6.0.3 + - Pygments==2.19.2 + - absl-py==2.1.0 + - accelerate==1.12.0 + - aiofiles==24.1.0 + - aiohappyeyeballs==2.6.1 + - aiohttp==3.13.3 + - aiosignal==1.4.0 + - annotated-doc==0.0.4 + - annotated-types==0.7.0 + - antlr4-python3-runtime==4.9.3 + - anyio==4.12.1 + - asciitree==0.3.3 + - asttokens==3.0.1 + - async-timeout==5.0.1 + - attrs==25.4.0 + - av==15.1.0 + - brotli==1.2.0 + - charset-normalizer==3.4.4 + - cmake==4.1.3 + - cmeel==0.58.0 + - cmeel-assimp==5.4.3.1 + - cmeel-boost==1.87.0.1 + - cmeel-console-bridge==1.0.2.3 + - cmeel-octomap==1.10.0 + - cmeel-qhull==8.0.2.1 + - cmeel-tinyxml==2.6.2.3 + - cmeel-tinyxml2==10.0.0 + - cmeel-urdfdom==3.1.1.1 + - cmeel-zlib==1.3.1 + - coal==3.0.2 + - coal-library==3.0.1 + - colorama==0.4.6 + - datasets==4.5.0 + - decorator==5.2.1 + - deepdiff==8.6.1 + - dill==0.4.0 + - docstring_parser==0.17.0 + - draccus==0.10.0 + - eigenpy==3.10.3 + - etils==1.7.0 + - evdev==1.9.2 + - exceptiongroup==1.3.1 + - executing==2.2.1 + - fastapi==0.128.0 + - fasteners==0.20 + - ffmpy==1.0.0 + - filelock==3.20.3 + - frozenlist==1.8.0 + - fsspec==2025.10.0 + - gitdb==4.0.12 + - glfw==2.7.0 + - gradio==6.3.0 + - gradio_client==2.0.3 + - groovy==0.1.2 + - gymnasium==1.2.3 + - h11==0.16.0 + - h5py==3.15.1 + - hf-xet==1.2.0 + - hf_transfer==0.1.9 + - httpcore==1.0.9 + - httpx==0.28.1 + - huggingface_hub==1.3.2 + - imageio==2.35.1 + - imageio-ffmpeg==0.6.0 + - importlib_metadata==8.7.1 + - importlib_resources==6.5.2 + - inquirerpy==0.3.4 + - ipython==8.38.0 + - jedi==0.19.2 + - jsonargparse==4.45.0 + - jsonlines==4.0.0 + - kiwisolver==1.4.5 + - lerobot==0.4.2 + - libcoal==3.0.2 + - libpinocchio==3.8.0 + - lightning==2.5.0.post0 + - lightning-utilities==0.15.2 + - lxml==5.3.0 + - markdown-it-py==4.0.0 + - matplotlib-inline==0.2.1 + - mdurl==0.1.2 + - mergedeep==1.3.4 + - mpmath==1.3.0 + - mujoco==3.2.2 + - mujoco-python-viewer==0.1.4 + - multidict==6.7.0 + - multiprocess==0.70.18 + - mypy_extensions==1.1.0 + - networkx==3.4.2 + - numcodecs==0.13.1 + - numpy==2.2.6 + - opencv-contrib-python==4.10.0.84 + - opencv-python==4.13.0.90 + - orderly-set==5.5.0 + - orjson==3.11.5 + - packaging==24.2 + - pandas==2.3.3 + - parso==0.8.5 + - pexpect==4.9.0 + - pfzy==0.3.4 + - pillow==12.1.0 + - pin==3.3.1 + - platformdirs==4.5.1 + - prompt_toolkit==3.0.52 + - propcache==0.4.1 + - protobuf==6.33.4 + - proxsuite==0.7.2 + - psutil==7.2.1 + - ptyprocess==0.7.0 + - pure_eval==0.2.3 + - pyarrow==22.0.0 + - pydantic==2.12.5 + - pydantic_core==2.41.5 + - pydub==0.25.1 + - pynput==1.8.1 + - pyquaternion==0.9.9 + - pyserial==3.5 + - python-dateutil==2.9.0.post0 + - python-multipart==0.0.21 + - python-xlib==0.33 + - pytorch-lightning==2.6.0 + - pyyaml-include==1.4.1 + - qwen-vl-utils==0.0.14 + - regex==2026.1.15 + - requests==2.32.5 + - rerun-sdk==0.26.2 + - rich==13.9.4 + - ruckig==0.9.2 + - safehttpx==0.1.7 + - safetensors==0.7.0 + - scipy==1.14.1 + - semantic-version==2.10.0 + - sentry-sdk==2.49.0 + - shellingham==1.5.4 + - smmap==5.0.2 + - stack-data==0.6.3 + - starlette==0.50.0 + - sympy==1.13.1 + - swanlab==0.7.13 + - termcolor==3.3.0 + - timm==1.0.24 + - toml==0.10.2 + - tomli==2.4.0 + - tomlkit==0.13.3 + - torchcodec==0.5 + - torchmetrics==1.8.2 + - tqdm==4.67.1 + - traitlets==5.14.3 + - typer==0.21.1 + - typer-slim==0.21.1 + - typeshed_client==2.8.2 + - typing-inspect==0.9.0 + - typing-inspection==0.4.2 + - typing_extensions==4.15.0 + - tzdata==2025.3 + - urdf_parser_py==0.0.4 + - urllib3==2.6.3 + - uv==0.9.28 + - uvicorn==0.40.0 + - wandb==0.24.0 + - wcwidth==0.2.14 + - xxhash==3.6.0 + - yarl==1.22.0 + - zarr==2.18.3 + - zipp==3.20.1 diff --git a/experiment_suites/2026-04-04-imf-horizon-grid/final_report.md b/experiment_suites/2026-04-04-imf-horizon-grid/final_report.md new file mode 100644 index 0000000..de0a215 --- /dev/null +++ b/experiment_suites/2026-04-04-imf-horizon-grid/final_report.md @@ -0,0 +1,63 @@ +# Phase-1 Final Report and Phase-2 Handoff + +- Finalized: 2026-04-05 00:34:20 CST +- Scope: IMF AttnRes policy horizon/action-step grid on `sim_transfer` +- Fixed setup: `n_emb=384`, `n_layer=12`, batch size `80`, learning rate `2.5e-4`, `max_steps=50k`, rollout every 5 epochs with 5 episodes, 3 cameras `[r_vis, top, front]`. +- Main metric: `checkpoints/vla_model_best.pt` 中记录的训练期最大 `rollout_avg_reward`。 + +## Final leaderboard + +| Rank | Run ID | pred_horizon | executed action steps | Best avg_reward | Best step | Final loss | +|---:|---|---:|---:|---:|---:|---:| +| 1 | `ph16_ex8` | 16 | 8 | **610.8** | 21874 | 0.0034 | +| 2 | `ph16_ex16` | 16 | 16 | 561.2 | 48124 | 0.0045 | +| 3 | `ph32_ex32` | 32 | 32 | 513.2 | 43749 | 0.0040 | +| 4 | `ph8_ex8` | 8 | 8 | 415.6 | 48124 | 0.0070 | +| 5 | `ph32_ex8` | 32 | 8 | 361.6 | 43749 | 0.0048 | +| 6 | `ph32_ex16` | 32 | 16 | 239.6 | 48124 | 0.0038 | + +## Final conclusions + +1. **最佳组合是 `pred_horizon=16` + `num_action_steps=8`**，最佳平均奖励为 **610.8**，出现在 **step 21874**。 +2. 在 `pred_horizon=16` 下，执行 8 步优于执行 16 步，优势约 **+8.8%**（610.8 vs 561.2）。 +3. `pred_horizon=32` 时，对执行步长非常敏感：`32/32` 明显优于 `32/8` 和 `32/16`；特别是 `32/16` 退化最明显。 +4. 更长的预测窗口并不会自动带来更高 reward；**预测窗口与实际执行窗口的匹配关系** 是关键。 +5. 最佳 checkpoint 并不在训练结束时出现，而是在 50k 训练中较早的 **21.9k step** 出现，说明 rollout 验证比仅看 train loss 更重要。 +6. 因而 Phase-2 的比较基线固定为 **`ph16_ex8`**。 + +## Recommended baseline for follow-up experiments + +- Baseline run: `ph16_ex8` +- Baseline best checkpoint: `step 21874` +- Baseline best avg_reward: `610.8` +- Baseline run dir: `/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223` + +## Phase-2 target: full-AttnRes vision backbone + +本阶段按你的要求，不再只是 IMF head 中使用 AttnRes，而是把**之前视觉 ResNet 主干中的残差单元全部替换为 AttnRes 残差单元**。当前实现保留了 ResNet 风格的 stage / downsample 宏观结构，但视觉残差 trunk 已切换到 AttnRes： + +- implementation: `roboimi/vla/models/backbones/attnres_resnet2d.py` +- wiring: `roboimi/vla/models/backbones/resnet_diffusion.py` +- config: `roboimi/vla/conf/backbone/resnet_diffusion.yaml` + +相关代码已提交： + +- `a780068` — headless rollout 修复 + Phase-1 汇总 +- `2033169` — full-AttnRes vision backbone + +## Phase-2 launch status (observed on 2026-04-05 00:36 CST) + +- Run: `imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-b40-lr1p25e4-ms50k-l20g3-20260405-002424` +- Host: `100.119.99.14`, GPU `3` +- Config anchor: `pred_horizon=16`, `num_action_steps=8` +- Vision backbone: `attnres_resnet` +- Because batch size `80` OOMed on both local 5090 and remote L20, Phase-2 currently uses: + - batch size: `40` + - learning rate: `1.25e-4` +- Latest confirmed progress: **step 1300** +- First rollout has **not happened yet** at this observation point. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/xy7fjdmn0stdr19eu3gub + +## Next action + +继续监控 Phase-2 full-AttnRes 训练，待其完成后直接与 Phase-1 baseline `610.8` 做对比，判断“视觉主干全部替换为 AttnRes”是否优于“仅 IMF 中使用 AttnRes”。 diff --git a/experiment_suites/2026-04-04-imf-horizon-grid/leaderboard.csv b/experiment_suites/2026-04-04-imf-horizon-grid/leaderboard.csv new file mode 100644 index 0000000..5031b26 --- /dev/null +++ b/experiment_suites/2026-04-04-imf-horizon-grid/leaderboard.csv @@ -0,0 +1,7 @@ +rank,run_id,status,pred_horizon,num_action_steps,best_rollout_avg_reward,best_step,final_step,final_loss,host,run_dir,latest_step +1,ph16_ex8,running,16,8,610.8,21874,50000,0.0034315965604037046,100.73.14.65,/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223,50000 +2,ph16_ex16,running,16,16,561.2,48124,50000,0.004544622730463743,100.119.99.14,/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223,50000 +3,ph32_ex32,finished,32,32,513.2,43749,50000,0.003953303210437298,local,/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223,49900 +4,ph8_ex8,running,8,8,415.6,48124,50000,0.007008877582848072,100.73.14.65,/home/droid/roboimi_suite_20260404/runs/imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223,50000 +5,ph32_ex8,running,32,8,361.6,43749,50000,0.004788532387465239,100.119.99.14,/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223,50000 +6,ph32_ex16,running,32,16,239.6,48124,50000,0.0038348555099219084,100.119.99.14,/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223,50000 diff --git a/experiment_suites/2026-04-04-imf-horizon-grid/manifest.json b/experiment_suites/2026-04-04-imf-horizon-grid/manifest.json new file mode 100644 index 0000000..862f384 --- /dev/null +++ b/experiment_suites/2026-04-04-imf-horizon-grid/manifest.json @@ -0,0 +1,115 @@ +{ + "suite_name": "2026-04-04-imf-horizon-grid", + "created_at": "2026-04-04 13:19:52", + "updated_at": "2026-04-04 13:19:52", + "phase": "phase1_launching", + "metric": "max_avg_reward", + "baseline": { + "agent": "resnet_imf_attnres", + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "val_split": 0.0, + "seed": 42, + "scheduler_type": "cosine", + "warmup_steps": 2000, + "min_lr": 1e-06, + "weight_decay": 1e-05, + "grad_clip": 1.0, + "inference_steps": 1, + "embed_dim": 384, + "n_layer": 12, + "n_head": 1, + "n_kv_head": 1, + "freeze_backbone": false, + "pretrained_backbone_weights": null, + "camera_names": [ + "r_vis", + "top", + "front" + ] + }, + "runs": [ + { + "id": "ph8_ex8", + "pred_horizon": 8, + "num_action_steps": 8, + "host": "100.73.14.65", + "host_label": "tailnet-5880", + "gpu": 0, + "workdir": "/home/droid/roboimi_suite_20260404", + "python": "/home/droid/miniforge3/envs/roboimi/bin/python", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "run_name": "imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223", + "launch_state": "ready" + }, + { + "id": "ph16_ex8", + "pred_horizon": 16, + "num_action_steps": 8, + "host": "100.73.14.65", + "host_label": "tailnet-5880", + "gpu": 1, + "workdir": "/home/droid/roboimi_suite_20260404", + "python": "/home/droid/miniforge3/envs/roboimi/bin/python", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "run_name": "imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223", + "launch_state": "ready" + }, + { + "id": "ph16_ex16", + "pred_horizon": 16, + "num_action_steps": 16, + "host": "100.119.99.14", + "host_label": "tailnet-l20", + "gpu": 0, + "workdir": "/home/droid/roboimi_suite_20260404", + "python": "/home/droid/miniforge3/envs/roboimi/bin/python", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "run_name": "imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223", + "launch_state": "provisioning_required" + }, + { + "id": "ph32_ex8", + "pred_horizon": 32, + "num_action_steps": 8, + "host": "100.119.99.14", + "host_label": "tailnet-l20", + "gpu": 1, + "workdir": "/home/droid/roboimi_suite_20260404", + "python": "/home/droid/miniforge3/envs/roboimi/bin/python", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "run_name": "imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223", + "launch_state": "provisioning_required" + }, + { + "id": "ph32_ex16", + "pred_horizon": 32, + "num_action_steps": 16, + "host": "100.119.99.14", + "host_label": "tailnet-l20", + "gpu": 2, + "workdir": "/home/droid/roboimi_suite_20260404", + "python": "/home/droid/miniforge3/envs/roboimi/bin/python", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "run_name": "imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223", + "launch_state": "provisioning_required" + }, + { + "id": "ph32_ex32", + "pred_horizon": 32, + "num_action_steps": 32, + "host": "local", + "host_label": "local-5090", + "gpu": 0, + "workdir": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy", + "python": "/home/droid/.conda/envs/roboimi/bin/python", + "dataset_dir": "/home/droid/project/diana_sim/sim_transfer", + "run_name": "imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223", + "launch_state": "ready" + } + ] +} diff --git a/experiment_suites/2026-04-04-imf-horizon-grid/notes.md b/experiment_suites/2026-04-04-imf-horizon-grid/notes.md new file mode 100644 index 0000000..e30da26 --- /dev/null +++ b/experiment_suites/2026-04-04-imf-horizon-grid/notes.md @@ -0,0 +1,20 @@ +# IMF Horizon Grid Suite Notes + +- Created: 2026-04-04 13:19:52 +- Phase-1 matrix: (8,8), (16,8), (16,16), (32,8), (32,16), (32,32) +- Fixed baseline: IMF AttnRes, n_emb=384, n_layer=12, batch_size=80, lr=2.5e-4, max_steps=50k, rollout every 5 epochs with 5 episodes. +- Host allocation: + - local RTX 5090: ph32_ex32 + - 100.73.14.65 RTX 5880 GPU0: ph8_ex8 + - 100.73.14.65 RTX 5880 GPU1: ph16_ex8 + - 100.119.99.14 L20 GPU0: ph16_ex16 + - 100.119.99.14 L20 GPU1: ph32_ex8 + - 100.119.99.14 L20 GPU2: ph32_ex16 +- 100.119.99.14 still needs env + dataset + swanlab credential copy before launch. + +- 2026-04-04 13:23:43: launched local ph32_ex32 (pid 1437836), remote 100.73 ph8_ex8 (pid 931824), ph16_ex8 (pid 931826); started 100.119 bootstrap (local pid 1437837). +- 2026-04-04 13:25:43: first status sync — local ph32_ex32 step≈500; remote ph8_ex8 step≈400; remote ph16_ex8 step≈400. +- 2026-04-04 13:27:41: second status sync — 100.119 bootstrap finished env copy and entered dataset copy; local ph32_ex32 step≈900; remote ph8_ex8 step≈800; remote ph16_ex8 step≈800. +- 2026-04-04 13:35:31: 100.119 bootstrap data/env copy finished. Original validation command hit a quoting bug, then I manually revalidated torch+mujoco+swanlab and launched ph16_ex16/ph32_ex8/ph32_ex16 with pids 81129/81130/81131. +- 2026-04-04 13:37:36: all 6 Phase-1 runs are now up. SwanLab links recorded in status.json; latest observed steps ~ local 900 / 5880 runs 800 / L20 runs 100. +- 2026-04-04 14:41:08: diagnosed remote first-rollout crash as early mujoco import before MUJOCO_GL=egl in eval_vla.py via raw_action_trajectory_viewer. Added regression test tests/test_eval_vla_headless_import.py, fixed import to lazy-load, verified 20-step headless eval on 5880 and L20, then resumed 5 failed runs from step 4374. Current resumed pids: ph8_ex8=938714, ph16_ex8=938717, ph16_ex16=90169, ph32_ex8=90173, ph32_ex16=90175. diff --git a/experiment_suites/2026-04-04-imf-horizon-grid/phase1_summary.md b/experiment_suites/2026-04-04-imf-horizon-grid/phase1_summary.md new file mode 100644 index 0000000..d8caafd --- /dev/null +++ b/experiment_suites/2026-04-04-imf-horizon-grid/phase1_summary.md @@ -0,0 +1,38 @@ +# Phase-1 IMF Horizon Grid Summary + +- Generated: 2026-04-04 23:43:38 +- Fixed baseline: IMF AttnRes head, n_emb=384, n_layer=12, batch_size=80, lr=2.5e-4, max_steps=50k, rollout every 5 epochs with 5 episodes, 3 cameras `[r_vis, top, front]`. +- Primary metric: `checkpoints/vla_model_best.pt -> rollout_avg_reward` (max training-time rollout average reward). + +## Ranked results + +| Rank | Run ID | pred_horizon | num_action_steps | Best avg_reward | Best step | Final loss | Host | +|---:|---|---:|---:|---:|---:|---:|---| +| 1 | `ph16_ex8` | 16 | 8 | 610.8 | 21874 | 0.0034 | 100.73.14.65 | +| 2 | `ph16_ex16` | 16 | 16 | 561.2 | 48124 | 0.0045 | 100.119.99.14 | +| 3 | `ph32_ex32` | 32 | 32 | 513.2 | 43749 | 0.0040 | local | +| 4 | `ph8_ex8` | 8 | 8 | 415.6 | 48124 | 0.0070 | 100.73.14.65 | +| 5 | `ph32_ex8` | 32 | 8 | 361.6 | 43749 | 0.0048 | 100.119.99.14 | +| 6 | `ph32_ex16` | 32 | 16 | 239.6 | 48124 | 0.0038 | 100.119.99.14 | + +## Main observations + +- Best overall setting was **`pred_horizon=16`, `num_action_steps=8`** with **max avg_reward = 610.8** at step **21874**. +- Comparing horizon 16: executing 8 steps outperformed executing 16 steps (`ph16_ex8` > `ph16_ex16`). +- Comparing horizon 32: executing the full 32-step chunk was much better than executing 16 or 8 steps (`ph32_ex32` > `ph32_ex8` > `ph32_ex16`). +- Short horizon 8 with 8-step execution was competitive but clearly below the best 16/8 and 32/32 settings. +- In this sweep, increasing prediction horizon helped only when the executed chunk length matched a good control cadence; mismatch could hurt a lot (especially `ph32_ex16`). + +## Raw results + +- `ph16_ex8`: best avg_reward=610.8 @ step 21874, final_loss=0.0034, run_dir=`/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223` +- `ph16_ex16`: best avg_reward=561.2 @ step 48124, final_loss=0.0045, run_dir=`/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223` +- `ph32_ex32`: best avg_reward=513.2 @ step 43749, final_loss=0.0040, run_dir=`/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223` +- `ph8_ex8`: best avg_reward=415.6 @ step 48124, final_loss=0.0070, run_dir=`/home/droid/roboimi_suite_20260404/runs/imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223` +- `ph32_ex8`: best avg_reward=361.6 @ step 43749, final_loss=0.0048, run_dir=`/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223` +- `ph32_ex16`: best avg_reward=239.6 @ step 48124, final_loss=0.0038, run_dir=`/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223` + +## Recommendation for Phase-2 anchor + +- Use **`pred_horizon=16`, `num_action_steps=8`** as the strongest Phase-1 baseline if the goal is purely maximizing rollout reward. +- If phase-2 needs a more conservative action execution budget, `ph16_ex8` is the strongest non-full-32 execution setting and may still be a good comparison anchor. diff --git a/experiment_suites/2026-04-04-imf-horizon-grid/status.json b/experiment_suites/2026-04-04-imf-horizon-grid/status.json new file mode 100644 index 0000000..e7f0bc8 --- /dev/null +++ b/experiment_suites/2026-04-04-imf-horizon-grid/status.json @@ -0,0 +1,167 @@ +{ + "suite_name": "2026-04-04-imf-horizon-grid", + "updated_at": "2026-04-05 00:34:20", + "phase": "phase1_completed", + "provisioning": { + "100.119.99.14": { + "state": "completed_manual_launch", + "bootstrap_pid_local": 1437837, + "log_path": "experiment_suites/2026-04-04-imf-horizon-grid/provision_logs/100.119.99.14-bootstrap-20260404-131223.log", + "env_copy": "completed", + "dataset_copy": "completed", + "launch_watcher_pid_local": null, + "launch_watcher_log": "experiment_suites/2026-04-04-imf-horizon-grid/launch_logs/100.119.99.14-launch-watcher-20260404-131223.log", + "swanlab_copy": "completed", + "bootstrap_validation_note": "initial validation command had a quoting bug; manual validation passed and launches were started successfully" + } + }, + "runs": { + "ph8_ex8": { + "status": "finished", + "host": "100.73.14.65", + "gpu": 0, + "run_name": "imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223", + "workdir": "/home/droid/roboimi_suite_20260404", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223/train_vla.log", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223", + "pred_horizon": 8, + "num_action_steps": 8, + "pid": 938714, + "launch_log": "experiment_suite_launch_logs/imf-p1-ph08-ex08-emb384-l12-ms50k-5880g0-20260404-131223.restartfix-20260404-143827.log", + "latest_step": 50000, + "latest_log_sync": "2026-04-05 00:34:20", + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/i5syc57b6zq7rbkrtqy7b", + "process_running": false, + "best_step": 48124, + "best_rollout_avg_reward": 415.6, + "final_loss": 0.007008877582848072 + }, + "ph16_ex8": { + "status": "finished", + "host": "100.73.14.65", + "gpu": 1, + "run_name": "imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223", + "workdir": "/home/droid/roboimi_suite_20260404", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223/train_vla.log", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223", + "pred_horizon": 16, + "num_action_steps": 8, + "pid": 938717, + "launch_log": "experiment_suite_launch_logs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223.restartfix-20260404-143827.log", + "latest_step": 50000, + "latest_log_sync": "2026-04-05 00:34:20", + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/4rusbrpfxmw4ffii1ul5w", + "process_running": false, + "best_step": 21874, + "best_rollout_avg_reward": 610.8, + "final_loss": 0.0034315965604037046 + }, + "ph16_ex16": { + "status": "finished", + "host": "100.119.99.14", + "gpu": 0, + "run_name": "imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223", + "workdir": "/home/droid/roboimi_suite_20260404", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223/train_vla.log", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223", + "pred_horizon": 16, + "num_action_steps": 16, + "pid": 90169, + "launch_log": "experiment_suite_launch_logs/imf-p1-ph16-ex16-emb384-l12-ms50k-l20g0-20260404-131223.restartfix-20260404-143827.log", + "latest_log_sync": "2026-04-05 00:34:20", + "latest_step": 50000, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/wwm232k6190gexnze8mg6", + "process_running": false, + "best_step": 48124, + "best_rollout_avg_reward": 561.2, + "final_loss": 0.004544622730463743 + }, + "ph32_ex8": { + "status": "finished", + "host": "100.119.99.14", + "gpu": 1, + "run_name": "imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223", + "workdir": "/home/droid/roboimi_suite_20260404", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223/train_vla.log", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223", + "pred_horizon": 32, + "num_action_steps": 8, + "pid": 90173, + "launch_log": "experiment_suite_launch_logs/imf-p1-ph32-ex08-emb384-l12-ms50k-l20g1-20260404-131223.restartfix-20260404-143827.log", + "latest_log_sync": "2026-04-05 00:34:20", + "latest_step": 50000, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/o5y2xjb2rsb3lmfcuhy4p", + "process_running": false, + "best_step": 43749, + "best_rollout_avg_reward": 361.6, + "final_loss": 0.004788532387465239 + }, + "ph32_ex16": { + "status": "finished", + "host": "100.119.99.14", + "gpu": 2, + "run_name": "imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223", + "workdir": "/home/droid/roboimi_suite_20260404", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223/train_vla.log", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223", + "pred_horizon": 32, + "num_action_steps": 16, + "pid": 90175, + "launch_log": "experiment_suite_launch_logs/imf-p1-ph32-ex16-emb384-l12-ms50k-l20g2-20260404-131223.restartfix-20260404-143827.log", + "latest_log_sync": "2026-04-05 00:34:20", + "latest_step": 50000, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/54cjpgba9eqsopdm0l8d3", + "process_running": false, + "best_step": 48124, + "best_rollout_avg_reward": 239.6, + "final_loss": 0.0038348555099219084 + }, + "ph32_ex32": { + "status": "finished", + "host": "local", + "gpu": 0, + "run_name": "imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223", + "workdir": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy", + "dataset_dir": "/home/droid/project/diana_sim/sim_transfer", + "log_path": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223/train_vla.log", + "run_dir": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223", + "pred_horizon": 32, + "num_action_steps": 32, + "pid": 1437836, + "launch_log": "experiment_suites/2026-04-04-imf-horizon-grid/launch_logs/imf-p1-ph32-ex32-emb384-l12-ms50k-5090-20260404-131223.launch.log", + "latest_step": 49900, + "latest_log_sync": "2026-04-05 00:34:20", + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/ajs2m218jd260hawhy5ns", + "process_running": false, + "latest_rollout_avg_reward": 513.2, + "best_rollout_avg_reward": 513.2, + "best_step": 43749, + "final_loss": 0.003953303210437298 + } + }, + "monitor": { + "state": "stopped", + "pid_local": null, + "log_path": "experiment_suites/2026-04-04-imf-horizon-grid/monitor_logs/status-sync-20260404-131223.log", + "interval_seconds": 300, + "stopped_at": "2026-04-05 00:34:20", + "stop_reason": "phase1 suite finalized after all six runs completed" + }, + "debug": { + "remote_rollout_failure_20260404": { + "root_cause": "eval_vla.py imported raw_action_trajectory_viewer at module import time, which imported mujoco before MUJOCO_GL=egl was set; remote headless rollout then fell back to GLFW/X11 and crashed with mujoco.FatalError: gladLoadGL error during env.reset()->mj.Renderer(...)", + "fixed_file": "roboimi/demos/vla_scripts/eval_vla.py", + "verification": { + "pytest": "tests/test_eval_vla_headless_import.py passed", + "remote_eval_5880": "1 episode x 20 steps headless eval passed", + "remote_eval_l20": "1 episode x 20 steps headless eval passed" + } + } + }, + "phase1_summary_md": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/experiment_suites/2026-04-04-imf-horizon-grid/phase1_summary.md" +} diff --git a/experiment_suites/2026-04-05-camera-ablation-summary.md b/experiment_suites/2026-04-05-camera-ablation-summary.md new file mode 100644 index 0000000..6b6d818 --- /dev/null +++ b/experiment_suites/2026-04-05-camera-ablation-summary.md @@ -0,0 +1,69 @@ +# Camera Ablation Summary (`pred_horizon=16`, `num_action_steps=8`, ResNet IMF) + +- Generated: 2026-04-05 +- Common setup: original ResNet vision backbone, `n_emb=384`, `n_layer=12`, `batch_size=80`, `lr=2.5e-4`, `max_steps=50k`, rollout every 5 epochs with 5 episodes, headless eval. +- Metric for comparison: `checkpoints/vla_model_best.pt -> rollout_avg_reward`. + +## Leaderboard + +| Rank | Cameras | Best avg_reward | Best step | Final loss | Run name | +|---:|---|---:|---:|---:|---| +| 1 | `top + front` | **274.8** | 48124 | 0.0056 | `imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023` | +| 2 | `top` | **271.2** | 43749 | 0.0052 | `imf-resnet-top-1cam-ph16-ex08-emb384-l12-ms50k-l20g4-20260405-125844` | +| 3 | `r_vis + front` | **244.0** | 21874 | 0.0043 | `imf-resnet-frontrvis-2cam-ph16-ex08-emb384-l12-ms50k-l20g1-20260405-102029` | +| 4 | `r_vis` | **6.4** | 17499 | 0.0047 | `imf-resnet-rvis-1cam-ph16-ex08-emb384-l12-ms50k-l20g3-20260405-125844` | +| 5 | `r_vis + top` | **1.2** | 4374 | 0.0047 | `imf-resnet-rvistop-2cam-ph16-ex08-emb384-l12-ms50k-l20g2-20260405-125844` | +| 6 | `front` | **0.0** | 4374 | 0.0074 | `imf-resnet-front-1cam-ph16-ex08-emb384-l12-ms50k-l20g0-20260405-095607` | + +## Main takeaways + +1. **`top` 是最关键的单相机视角**：`top only = 271.2`，几乎与 `top + front = 274.8` 持平。 +2. **`front` 单独几乎没有效用**：`front only = 0.0`。 +3. **`r_vis` 单独也基本无效**：`r_vis only = 6.4`。 +4. **`r_vis + front` 可以显著优于单独 `front` / `r_vis`**，说明这两个视角有一定互补性，但仍明显弱于任何包含 `top` 且表现正常的配置。 +5. **`r_vis + top` 的结果异常差**：只有 `1.2`，远低于 `top only = 271.2`。这说明简单加入 `r_vis` 并不保证增益，甚至可能破坏当前设置下的学习。 +6. **训练 loss 与 rollout reward 明显不一致**：例如 `r_vis + top` 和 `r_vis only` 的 final loss 都不高，但 reward 很差，因此本组实验必须以 rollout reward 而不是 loss 选型。 + +## Horizontal comparison views + +### Single-camera comparison + +- `top`: **271.2** +- `r_vis`: **6.4** +- `front`: **0.0** + +结论：**`top >>> r_vis > front`**。 + +### Two-camera comparison + +- `top + front`: **274.8** +- `r_vis + front`: **244.0** +- `r_vis + top`: **1.2** + +结论： +- **最稳妥的双相机组合是 `top + front`**。 +- `r_vis + front` 有效，但不如 `top + front`。 +- `r_vis + top` 在当前设置下几乎失效。 + +### Incremental effect of adding a second view + +- 在 `top` 基础上加 `front`：`271.2 -> 274.8`，**增益很小**。 +- 在 `front` 基础上加 `r_vis`：`0.0 -> 244.0`，**增益很大**。 +- 在 `top` 基础上加 `r_vis`：`271.2 -> 1.2`，**显著退化**。 + +## Practical recommendation + +如果只从这 6 个实验里选： + +- **首选**：`top + front` +- **次选**：`top only` +- 如果必须不用 `top`：`r_vis + front` 明显优于 `front only` / `r_vis only` +- **不建议**：`r_vis + top` + +## Note relative to previous 3-camera baseline + +此前 3 相机 `[r_vis, top, front]` 的最佳 reward 为 **610.8**。 +因此这次 6 个 camera ablation 的最佳结果（`top + front = 274.8`）说明： + +- 当前这个训练批次里，**去掉任意一个视角都会显著低于之前的 3 相机最优结果**； +- 但在去掉视角的约束下，**`top` 仍然是最核心的保留对象**。 diff --git a/experiment_suites/2026-04-05-front-only-resnet-1cam/CHECKLIST.md b/experiment_suites/2026-04-05-front-only-resnet-1cam/CHECKLIST.md new file mode 100644 index 0000000..7999908 --- /dev/null +++ b/experiment_suites/2026-04-05-front-only-resnet-1cam/CHECKLIST.md @@ -0,0 +1,8 @@ +# CHECKLIST + +- [x] Confirm remote free GPU +- [x] Create front-only run contract +- [x] Remote smoke test passes +- [x] Launch 50k run on remote GPU0 +- [x] Record pid / log / SwanLab +- [x] Report status back to user diff --git a/experiment_suites/2026-04-05-front-only-resnet-1cam/PLAN.md b/experiment_suites/2026-04-05-front-only-resnet-1cam/PLAN.md new file mode 100644 index 0000000..4c8ea2d --- /dev/null +++ b/experiment_suites/2026-04-05-front-only-resnet-1cam/PLAN.md @@ -0,0 +1,28 @@ +# PLAN + +## Goal +Train a 50k-step IMF baseline with the original ResNet vision backbone, using only the `front` camera as image conditioning. + +## Fixed comparison contract +- Same as the active `top/front` run except image input is reduced to `[front]` +- Agent: `resnet_imf_attnres` +- Vision backbone mode: `resnet` +- `pred_horizon=16`, `num_action_steps=8` +- `n_emb=384`, `n_layer=12`, `n_head=1`, `n_kv_head=1` +- `inference_steps=1` +- `batch_size=80`, `lr=2.5e-4`, cosine, warmup=2000 +- dataset: `/home/droid/sim_dataset/sim_transfer` +- cameras: `[front]` only +- rollout every 5 epochs with 5 episodes, headless + +## Resource plan +- Host: `100.119.99.14` +- GPU: `0` + +## Important dimension override +- Single-camera visual cond dim = `64 + 16 = 80`, so override `agent.head.cond_dim=80` and `agent.num_cams=1`. + +## Execution path +1. 2-step smoke test on remote GPU0. +2. If smoke passes, launch 50k main run with SwanLab. +3. Record pid / run_dir / log / URL locally. diff --git a/experiment_suites/2026-04-05-front-only-resnet-1cam/notes.md b/experiment_suites/2026-04-05-front-only-resnet-1cam/notes.md new file mode 100644 index 0000000..212db50 --- /dev/null +++ b/experiment_suites/2026-04-05-front-only-resnet-1cam/notes.md @@ -0,0 +1,6 @@ +# Notes + +- 2026-04-05 09:55:27: remote 2-step smoke passed on `100.119.99.14` GPU0 with `front` only, batch=80, no OOM. +- 2026-04-05 09:56:26: launched main run `imf-resnet-front-1cam-ph16-ex08-emb384-l12-ms50k-l20g0-20260405-095607`. +- 2026-04-05 09:57:36: confirmed training is stable through step 200, latest loss 0.2830. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/7kdii8oc6tjkcyu5y0lwq diff --git a/experiment_suites/2026-04-05-front-only-resnet-1cam/status.json b/experiment_suites/2026-04-05-front-only-resnet-1cam/status.json new file mode 100644 index 0000000..77aaa80 --- /dev/null +++ b/experiment_suites/2026-04-05-front-only-resnet-1cam/status.json @@ -0,0 +1,51 @@ +{ + "suite_name": "2026-04-05-front-only-resnet-1cam", + "updated_at": "2026-04-05 09:57:36", + "phase": "running", + "baseline_reference": { + "source_run": "imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023", + "notes": "Same hyperparameters as the active top/front run, but image input is reduced to [front] only." + }, + "smoke_test": { + "status": "passed", + "host": "100.119.99.14", + "gpu": 0, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/smoke-frontonly-resnet-ph16-ex08-20260405-095509", + "batch_size": 80, + "max_steps": 2, + "note": "2-step remote CUDA smoke passed on L20 GPU0 without OOM." + }, + "main_run": { + "status": "running", + "host": "100.119.99.14", + "gpu": 0, + "launch_pid": 158874, + "pid": 158877, + "run_name": "imf-resnet-front-1cam-ph16-ex08-emb384-l12-ms50k-l20g0-20260405-095607", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-front-1cam-ph16-ex08-emb384-l12-ms50k-l20g0-20260405-095607", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-front-1cam-ph16-ex08-emb384-l12-ms50k-l20g0-20260405-095607/train_vla.log", + "launch_log": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/imf-resnet-front-1cam-ph16-ex08-emb384-l12-ms50k-l20g0-20260405-095607.launch.log", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "camera_names": [ + "front" + ], + "pred_horizon": 16, + "num_action_steps": 8, + "head_cond_dim": 80, + "head_n_emb": 384, + "head_n_layer": 12, + "vision_backbone_mode": "resnet", + "pretrained_backbone_weights": null, + "freeze_backbone": false, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/7kdii8oc6tjkcyu5y0lwq", + "latest_step": 200, + "latest_loss": 0.283, + "process_running": true + } +} diff --git a/experiment_suites/2026-04-05-front-rvis-resnet-2cam/CHECKLIST.md b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/CHECKLIST.md new file mode 100644 index 0000000..c5a5c9e --- /dev/null +++ b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/CHECKLIST.md @@ -0,0 +1,8 @@ +# CHECKLIST + +- [x] Confirm camera mapping (`right` -> `r_vis`) +- [x] Create front+r_vis run contract +- [x] Remote smoke test passes +- [x] Launch 50k run on remote GPU1 +- [x] Record pid / log / SwanLab +- [x] Report status back to user diff --git a/experiment_suites/2026-04-05-front-rvis-resnet-2cam/PLAN.md b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/PLAN.md new file mode 100644 index 0000000..82f7d1b --- /dev/null +++ b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/PLAN.md @@ -0,0 +1,23 @@ +# PLAN + +## Goal +Train a 50k-step IMF baseline with the original ResNet vision backbone, using `front` + `r_vis` cameras only. + +## Fixed comparison contract +- Same hyperparameters as the active top/front and front-only runs +- Agent: `resnet_imf_attnres` +- Vision backbone mode: `resnet` +- `pred_horizon=16`, `num_action_steps=8` +- `n_emb=384`, `n_layer=12`, `n_head=1`, `n_kv_head=1` +- `inference_steps=1` +- `batch_size=80`, `lr=2.5e-4`, cosine warmup 2000 +- dataset: `/home/droid/sim_dataset/sim_transfer` +- cameras: `[r_vis, front]` +- rollout every 5 epochs with 5 episodes, headless + +## Important dimension override +- Two-camera visual cond dim = `64*2 + 16 = 144`, so set `agent.num_cams=2`, `agent.head.cond_dim=144`. + +## Resource plan +- Host: `100.119.99.14` +- GPU: `1` diff --git a/experiment_suites/2026-04-05-front-rvis-resnet-2cam/notes.md b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/notes.md new file mode 100644 index 0000000..3699a46 --- /dev/null +++ b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/notes.md @@ -0,0 +1,6 @@ +# Notes + +- 2026-04-05 10:20:09: remote 2-step smoke passed on `100.119.99.14` GPU1 with `r_vis + front`, batch=80, no OOM. +- 2026-04-05 10:20:49: launched main run `imf-resnet-frontrvis-2cam-ph16-ex08-emb384-l12-ms50k-l20g1-20260405-102029`. +- 2026-04-05 10:22:03: confirmed training is stable through step 200, latest loss 0.3321. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/3fyzjfdcbiq7frtbqv6ss diff --git a/experiment_suites/2026-04-05-front-rvis-resnet-2cam/status.json b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/status.json new file mode 100644 index 0000000..4f17c98 --- /dev/null +++ b/experiment_suites/2026-04-05-front-rvis-resnet-2cam/status.json @@ -0,0 +1,55 @@ +{ + "suite_name": "2026-04-05-front-rvis-resnet-2cam", + "updated_at": "2026-04-05 10:22:03", + "phase": "running", + "interpretation": { + "right_camera_name": "r_vis" + }, + "baseline_reference": { + "source_run": "imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023", + "notes": "Same hyperparameters as the active top/front run, replacing top with r_vis." + }, + "smoke_test": { + "status": "passed", + "host": "100.119.99.14", + "gpu": 1, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/smoke-frontrvis-resnet-ph16-ex08-20260405-102001", + "batch_size": 80, + "max_steps": 2, + "note": "2-step remote CUDA smoke passed on L20 GPU1 without OOM." + }, + "main_run": { + "status": "running", + "host": "100.119.99.14", + "gpu": 1, + "launch_pid": 159910, + "pid": 159913, + "run_name": "imf-resnet-frontrvis-2cam-ph16-ex08-emb384-l12-ms50k-l20g1-20260405-102029", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-frontrvis-2cam-ph16-ex08-emb384-l12-ms50k-l20g1-20260405-102029", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-frontrvis-2cam-ph16-ex08-emb384-l12-ms50k-l20g1-20260405-102029/train_vla.log", + "launch_log": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/imf-resnet-frontrvis-2cam-ph16-ex08-emb384-l12-ms50k-l20g1-20260405-102029.launch.log", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "camera_names": [ + "r_vis", + "front" + ], + "pred_horizon": 16, + "num_action_steps": 8, + "head_cond_dim": 144, + "head_n_emb": 384, + "head_n_layer": 12, + "vision_backbone_mode": "resnet", + "pretrained_backbone_weights": null, + "freeze_backbone": false, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/3fyzjfdcbiq7frtbqv6ss", + "latest_step": 200, + "latest_loss": 0.3321, + "process_running": true + } +} diff --git a/experiment_suites/2026-04-05-full-attnres-vision-phase2/manifest.json b/experiment_suites/2026-04-05-full-attnres-vision-phase2/manifest.json new file mode 100644 index 0000000..2be19e9 --- /dev/null +++ b/experiment_suites/2026-04-05-full-attnres-vision-phase2/manifest.json @@ -0,0 +1,20 @@ +{ + "suite_name": "2026-04-05-full-attnres-vision-phase2", + "created_at": "2026-04-05 00:12:14", + "phase": "phase2_running", + "baseline_reference": { + "run_id": "ph16_ex8", + "best_rollout_avg_reward": 610.8, + "best_step": 21874, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223" + }, + "candidate": { + "run_name": "imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-ms50k-20260405-001214", + "host": "local", + "gpu": 0, + "pred_horizon": 16, + "num_action_steps": 8, + "vision_backbone_mode": "attnres_resnet", + "notes": "Full-AttnRes vision backbone replacing ResNet residual units; IMF head unchanged." + } +} diff --git a/experiment_suites/2026-04-05-full-attnres-vision-phase2/notes.md b/experiment_suites/2026-04-05-full-attnres-vision-phase2/notes.md new file mode 100644 index 0000000..d6bb4dd --- /dev/null +++ b/experiment_suites/2026-04-05-full-attnres-vision-phase2/notes.md @@ -0,0 +1,9 @@ +# Full-AttnRes Vision Phase-2 + +- Created: 2026-04-05 00:12:14 +- Baseline reference: ph16_ex8 best avg_reward=610.8 +- Candidate run: imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-ms50k-20260405-001214 +- 2026-04-05 00:23:03: batch=80 OOM on both 5090 and L20; using validated fallback batch=40, lr=1.25e-4 on remote L20 GPU3. +- 2026-04-05 00:24:24: launching candidate imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-b40-lr1p25e4-ms50k-l20g3-20260405-002424 on 100.119.99.14 GPU3 with batch=40 lr=1.25e-4. +- 2026-04-05 00:27:17: remote phase2 run is active on 100.119.99.14 GPU3, validated at least to step 200. SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/xy7fjdmn0stdr19eu3gub +- 2026-04-05 00:36:54: latest confirmed progress is step 1300 on 100.119.99.14 GPU3; first rollout not reached yet. diff --git a/experiment_suites/2026-04-05-full-attnres-vision-phase2/status.json b/experiment_suites/2026-04-05-full-attnres-vision-phase2/status.json new file mode 100644 index 0000000..18b855f --- /dev/null +++ b/experiment_suites/2026-04-05-full-attnres-vision-phase2/status.json @@ -0,0 +1,32 @@ +{ + "suite_name": "2026-04-05-full-attnres-vision-phase2", + "updated_at": "2026-04-05 00:36:54", + "phase": "phase2_running", + "baseline_reference": { + "run_id": "ph16_ex8", + "best_rollout_avg_reward": 610.8, + "best_step": 21874, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223" + }, + "candidate": { + "run_name": "imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-b40-lr1p25e4-ms50k-l20g3-20260405-002424", + "host": "100.119.99.14", + "gpu": 3, + "pred_horizon": 16, + "num_action_steps": 8, + "vision_backbone_mode": "attnres_resnet", + "notes": "Full-AttnRes vision backbone replacing ResNet residual units; IMF head unchanged.", + "status": "running", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-b40-lr1p25e4-ms50k-l20g3-20260405-002424", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-b40-lr1p25e4-ms50k-l20g3-20260405-002424/train_vla.log", + "pid": 151187, + "batch_size": 40, + "lr": 0.000125, + "num_workers": 12, + "launch_log": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/imf-p2-full-attnres-vision-ph16-ex08-emb384-l12-b40-lr1p25e4-ms50k-l20g3-20260405-002424.launch.log", + "note": "Local 5090 and remote L20 both OOM at batch=80; switched to batch=40 and linearly scaled lr to 1.25e-4 after smoke validation on L20.", + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/xy7fjdmn0stdr19eu3gub", + "latest_step": 1300, + "latest_log_sync": "2026-04-05 00:36:54" + } +} diff --git a/experiment_suites/2026-04-05-lewm-vit-transfer/manifest.json b/experiment_suites/2026-04-05-lewm-vit-transfer/manifest.json new file mode 100644 index 0000000..4f44477 --- /dev/null +++ b/experiment_suites/2026-04-05-lewm-vit-transfer/manifest.json @@ -0,0 +1,73 @@ +{ + "date": "2026-04-06", + "branch": "feat-imf-attnres-policy", + "worktree": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy", + "model": "LEWM ViT frozen visual encoder + IMF AttnRes diffusion head", + "checkpoint_path": "/home/droid/le-wm/lewm-sim-transfer/pa1w85md8jop6bvol8oxp/checkpoints/epoch=99-step=47800.ckpt", + "visual_contract": { + "input_camera_names": ["r_vis", "top", "front"], + "fused_camera_names": ["front", "top", "r_vis"], + "joint_output_dim": 192, + "freeze_backbone": true, + "dataset_image_resize_shape": null, + "eval_image_resize_shape": [256, 256], + "fused_short_side_resize": 224 + }, + "training_contract": { + "pred_horizon": 16, + "num_action_steps": 8, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 10, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "scheduler_type": "cosine", + "warmup_steps": 2000, + "min_lr": 1e-06, + "weight_decay": 1e-05, + "grad_clip": 1.0 + }, + "verification": { + "local_tests": "38 passed", + "remote_dataset_shape": [2, 3, 256, 256], + "remote_eval_prepared_shape": [3, 256, 256], + "remote_smoke_run": { + "run_name": "smoke-lewm-imf-rawpath-emb384-20260406-002002", + "result": "passed", + "details": "2-step train + checkpoint-triggered 1-episode headless rollout succeeded with corrected raw256 path" + } + }, + "superseded_runs": [ + { + "run_name": "lewm-vit-imf-sim-transfer-emb384-l12-ph16-ex08-step50k-roll10-5880g0-20260405-201914", + "reason": "stopped due to incorrect early per-camera 224 resize" + }, + { + "run_name": "lewm-vit-imf-sim-transfer-emb256-l12-ph16-ex08-step50k-roll10-5880g1-20260405-201914", + "reason": "stopped due to incorrect early per-camera 224 resize" + } + ], + "full_runs": [ + { + "host": "100.73.14.65", + "gpu": 0, + "run_name": "lewm-vit-imf-raw256fix-sim-transfer-emb384-l12-ph16-ex08-step50k-roll10-5880g0-20260406-002124", + "pid": 1058589, + "log_path": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/lewm-vit-imf-raw256fix-sim-transfer-emb384-l12-ph16-ex08-step50k-roll10-5880g0-20260406-002124.launch.log", + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/y5tzgqe0u966w9ak41i31", + "head_n_emb": 384, + "head_n_layer": 12 + }, + { + "host": "100.73.14.65", + "gpu": 1, + "run_name": "lewm-vit-imf-raw256fix-sim-transfer-emb256-l12-ph16-ex08-step50k-roll10-5880g1-20260406-002124", + "pid": 1058590, + "log_path": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/lewm-vit-imf-raw256fix-sim-transfer-emb256-l12-ph16-ex08-step50k-roll10-5880g1-20260406-002124.launch.log", + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/2esr9y7t2dgesstgrn5i6", + "head_n_emb": 256, + "head_n_layer": 12 + } + ] +} diff --git a/experiment_suites/2026-04-05-lewm-vit-transfer/notes.md b/experiment_suites/2026-04-05-lewm-vit-transfer/notes.md new file mode 100644 index 0000000..fbde1c4 --- /dev/null +++ b/experiment_suites/2026-04-05-lewm-vit-transfer/notes.md @@ -0,0 +1,25 @@ +# 2026-04-06 LEWM ViT Transfer Notes + +## Root-cause fix + +The first LEWM runs were stopped because the data path still resized each camera view to `224x224` **before** multiview fusion. That preserved the final tensor shape but broke the original LEWM geometry. + +Corrected path now is: + +- **Training dataset**: keep stored per-view `256x256` images (`data.image_resize_shape=null` at launch; dataset instantiate override is `None` for LEWM) +- **Eval rollout input**: resize live MuJoCo `480x640` camera images to `256x256` per view +- **Backbone**: fuse `front, top, r_vis` on the LEWM axis, then resize fused short side to `224` + +## Verification + +- Local tests passed (`38 passed` across the focused suite) +- Remote check: + - dataset sample image shape: `(2, 3, 256, 256)` + - eval-prepared live frame shape: `(3, 256, 256)` +- Remote smoke passed with real checkpoint: + - `smoke-lewm-imf-rawpath-emb384-20260406-002002` + +## Current runs + +- `lewm-vit-imf-raw256fix-sim-transfer-emb384-l12-ph16-ex08-step50k-roll10-5880g0-20260406-002124` +- `lewm-vit-imf-raw256fix-sim-transfer-emb256-l12-ph16-ex08-step50k-roll10-5880g1-20260406-002124` diff --git a/experiment_suites/2026-04-05-lewm-vit-transfer/status.json b/experiment_suites/2026-04-05-lewm-vit-transfer/status.json new file mode 100644 index 0000000..b27c604 --- /dev/null +++ b/experiment_suites/2026-04-05-lewm-vit-transfer/status.json @@ -0,0 +1,19 @@ +{ + "status": "running", + "updated_at": "2026-04-06T00:22:10+08:00", + "remote_host": "100.73.14.65", + "runs": [ + { + "run_name": "lewm-vit-imf-raw256fix-sim-transfer-emb384-l12-ph16-ex08-step50k-roll10-5880g0-20260406-002124", + "pid": 1058589, + "gpu": 0, + "state": "running" + }, + { + "run_name": "lewm-vit-imf-raw256fix-sim-transfer-emb256-l12-ph16-ex08-step50k-roll10-5880g1-20260406-002124", + "pid": 1058590, + "gpu": 1, + "state": "running" + } + ] +} diff --git a/experiment_suites/2026-04-05-rvis-only-resnet-1cam/CHECKLIST.md b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/CHECKLIST.md new file mode 100644 index 0000000..764a396 --- /dev/null +++ b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/CHECKLIST.md @@ -0,0 +1,7 @@ +# CHECKLIST + +- [x] Create run contract +- [x] Remote smoke test passes +- [x] Launch 50k main run +- [x] Record pid / log / SwanLab +- [x] Report status back to user diff --git a/experiment_suites/2026-04-05-rvis-only-resnet-1cam/PLAN.md b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/PLAN.md new file mode 100644 index 0000000..f473389 --- /dev/null +++ b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/PLAN.md @@ -0,0 +1,12 @@ +# PLAN + +## Goal +Train a 50k-step IMF baseline with the original ResNet vision backbone using r_vis only as the only image conditioning. + +## Fixed comparison contract +- same hyperparameters as the active top/front run +- cameras: ['r_vis'] +- num_cams=1 +- head.cond_dim=80 +- host: 100.119.99.14 +- gpu: 3 diff --git a/experiment_suites/2026-04-05-rvis-only-resnet-1cam/notes.md b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/notes.md new file mode 100644 index 0000000..546006a --- /dev/null +++ b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/notes.md @@ -0,0 +1,6 @@ +# Notes + +- 2026-04-05 12:58:22: smoke passed for ['r_vis'] on 100.119.99.14 GPU3. +- 2026-04-05 12:59:24: launched main run `imf-resnet-rvis-1cam-ph16-ex08-emb384-l12-ms50k-l20g3-20260405-125844`. +- 2026-04-05 13:01:20: latest confirmed progress step=400, loss=0.1165. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/qnuh7vln9mqomxxldyecq diff --git a/experiment_suites/2026-04-05-rvis-only-resnet-1cam/status.json b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/status.json new file mode 100644 index 0000000..8fbe1ca --- /dev/null +++ b/experiment_suites/2026-04-05-rvis-only-resnet-1cam/status.json @@ -0,0 +1,47 @@ +{ + "suite_name": "2026-04-05-rvis-only-resnet-1cam", + "updated_at": "2026-04-05 13:01:20", + "phase": "running", + "smoke_test": { + "status": "passed", + "host": "100.119.99.14", + "gpu": 3, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/smoke-rvisonly-resnet-ph16-ex08-20260405-125812", + "batch_size": 80, + "max_steps": 2, + "note": "2-step remote CUDA smoke passed without OOM." + }, + "main_run": { + "status": "running", + "host": "100.119.99.14", + "gpu": 3, + "launch_pid": 164812, + "pid": 164816, + "run_name": "imf-resnet-rvis-1cam-ph16-ex08-emb384-l12-ms50k-l20g3-20260405-125844", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-rvis-1cam-ph16-ex08-emb384-l12-ms50k-l20g3-20260405-125844", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-rvis-1cam-ph16-ex08-emb384-l12-ms50k-l20g3-20260405-125844/train_vla.log", + "launch_log": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/imf-resnet-rvis-1cam-ph16-ex08-emb384-l12-ms50k-l20g3-20260405-125844.launch.log", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "camera_names": [ + "r_vis" + ], + "pred_horizon": 16, + "num_action_steps": 8, + "head_cond_dim": 80, + "head_n_emb": 384, + "head_n_layer": 12, + "vision_backbone_mode": "resnet", + "pretrained_backbone_weights": null, + "freeze_backbone": false, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/qnuh7vln9mqomxxldyecq", + "latest_step": 400, + "latest_loss": 0.1165, + "process_running": true + } +} diff --git a/experiment_suites/2026-04-05-rvistop-resnet-2cam/CHECKLIST.md b/experiment_suites/2026-04-05-rvistop-resnet-2cam/CHECKLIST.md new file mode 100644 index 0000000..764a396 --- /dev/null +++ b/experiment_suites/2026-04-05-rvistop-resnet-2cam/CHECKLIST.md @@ -0,0 +1,7 @@ +# CHECKLIST + +- [x] Create run contract +- [x] Remote smoke test passes +- [x] Launch 50k main run +- [x] Record pid / log / SwanLab +- [x] Report status back to user diff --git a/experiment_suites/2026-04-05-rvistop-resnet-2cam/PLAN.md b/experiment_suites/2026-04-05-rvistop-resnet-2cam/PLAN.md new file mode 100644 index 0000000..459d2c5 --- /dev/null +++ b/experiment_suites/2026-04-05-rvistop-resnet-2cam/PLAN.md @@ -0,0 +1,12 @@ +# PLAN + +## Goal +Train a 50k-step IMF baseline with the original ResNet vision backbone using r_vis + top as the only image conditioning. + +## Fixed comparison contract +- same hyperparameters as the active top/front run +- cameras: ['r_vis', 'top'] +- num_cams=2 +- head.cond_dim=144 +- host: 100.119.99.14 +- gpu: 2 diff --git a/experiment_suites/2026-04-05-rvistop-resnet-2cam/notes.md b/experiment_suites/2026-04-05-rvistop-resnet-2cam/notes.md new file mode 100644 index 0000000..02767af --- /dev/null +++ b/experiment_suites/2026-04-05-rvistop-resnet-2cam/notes.md @@ -0,0 +1,6 @@ +# Notes + +- 2026-04-05 12:58:22: smoke passed for ['r_vis', 'top'] on 100.119.99.14 GPU2. +- 2026-04-05 12:59:24: launched main run `imf-resnet-rvistop-2cam-ph16-ex08-emb384-l12-ms50k-l20g2-20260405-125844`. +- 2026-04-05 13:01:20: latest confirmed progress step=200, loss=0.2845. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/umsm6402eb81et7wx7z4a diff --git a/experiment_suites/2026-04-05-rvistop-resnet-2cam/status.json b/experiment_suites/2026-04-05-rvistop-resnet-2cam/status.json new file mode 100644 index 0000000..525d542 --- /dev/null +++ b/experiment_suites/2026-04-05-rvistop-resnet-2cam/status.json @@ -0,0 +1,48 @@ +{ + "suite_name": "2026-04-05-rvistop-resnet-2cam", + "updated_at": "2026-04-05 13:01:20", + "phase": "running", + "smoke_test": { + "status": "passed", + "host": "100.119.99.14", + "gpu": 2, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/smoke-rvistop-resnet-ph16-ex08-20260405-125812", + "batch_size": 80, + "max_steps": 2, + "note": "2-step remote CUDA smoke passed without OOM." + }, + "main_run": { + "status": "running", + "host": "100.119.99.14", + "gpu": 2, + "launch_pid": 164745, + "pid": 164749, + "run_name": "imf-resnet-rvistop-2cam-ph16-ex08-emb384-l12-ms50k-l20g2-20260405-125844", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-rvistop-2cam-ph16-ex08-emb384-l12-ms50k-l20g2-20260405-125844", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-rvistop-2cam-ph16-ex08-emb384-l12-ms50k-l20g2-20260405-125844/train_vla.log", + "launch_log": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/imf-resnet-rvistop-2cam-ph16-ex08-emb384-l12-ms50k-l20g2-20260405-125844.launch.log", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "camera_names": [ + "r_vis", + "top" + ], + "pred_horizon": 16, + "num_action_steps": 8, + "head_cond_dim": 144, + "head_n_emb": 384, + "head_n_layer": 12, + "vision_backbone_mode": "resnet", + "pretrained_backbone_weights": null, + "freeze_backbone": false, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/umsm6402eb81et7wx7z4a", + "latest_step": 200, + "latest_loss": 0.2845, + "process_running": true + } +} diff --git a/experiment_suites/2026-04-05-top-front-resnet-2cam/CHECKLIST.md b/experiment_suites/2026-04-05-top-front-resnet-2cam/CHECKLIST.md new file mode 100644 index 0000000..2ef9cad --- /dev/null +++ b/experiment_suites/2026-04-05-top-front-resnet-2cam/CHECKLIST.md @@ -0,0 +1,8 @@ +# CHECKLIST + +- [x] Confirm baseline hyperparameters from trusted prior run +- [x] Confirm local GPU availability +- [x] Smoke test with `top/front` cameras only +- [x] Launch 50k run +- [x] Record pid / run dir / log path / SwanLab URL +- [x] Report status back to user diff --git a/experiment_suites/2026-04-05-top-front-resnet-2cam/PLAN.md b/experiment_suites/2026-04-05-top-front-resnet-2cam/PLAN.md new file mode 100644 index 0000000..68ff490 --- /dev/null +++ b/experiment_suites/2026-04-05-top-front-resnet-2cam/PLAN.md @@ -0,0 +1,30 @@ +# PLAN + +## Goal +Train a 50k-step IMF baseline with the original ResNet vision backbone (no full-AttnRes vision replacement), using only `top` and `front` cameras as image conditioning. + +## Fixed comparison contract +- Agent: `resnet_imf_attnres` +- Vision backbone mode: `resnet` +- `pred_horizon=16` +- `num_action_steps=8` +- `n_emb=384`, `n_layer=12`, `n_head=1`, `n_kv_head=1` +- `inference_steps=1` +- `batch_size=80`, `lr=2.5e-4`, cosine scheduler, warmup 2000 +- dataset: `/home/droid/project/diana_sim/sim_transfer` +- cameras: `[top, front]` only +- training budget: `max_steps=50000` +- rollout validation: every 5 epochs, 5 episodes, headless + +## Resource plan +- Host: local +- GPU: RTX 5090 (GPU 0) + +## Execution path +1. Run a short 2-step smoke test on GPU with the exact 2-camera config. +2. If smoke passes, launch the 50k main run with durable log redirection. +3. Record run name, pid, log path, and SwanLab URL into suite status. + +## Fallbacks +- If batch 80 OOMs, fall back to batch 64 with scaled lr 2.0e-4. +- If dataloader startup is unstable, reduce num_workers from 12 to 8. diff --git a/experiment_suites/2026-04-05-top-front-resnet-2cam/notes.md b/experiment_suites/2026-04-05-top-front-resnet-2cam/notes.md new file mode 100644 index 0000000..d859f77 --- /dev/null +++ b/experiment_suites/2026-04-05-top-front-resnet-2cam/notes.md @@ -0,0 +1,5 @@ +# Notes + +- 2026-04-05 08:50:04: 2-step smoke test passed locally on RTX 5090 with `top/front` cameras, batch=80, no OOM. +- 2026-04-05 08:50:42: launched main run `imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023` on local GPU0. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/vi77mn5dwd19z4nttxab8 diff --git a/experiment_suites/2026-04-05-top-front-resnet-2cam/status.json b/experiment_suites/2026-04-05-top-front-resnet-2cam/status.json new file mode 100644 index 0000000..1a88b47 --- /dev/null +++ b/experiment_suites/2026-04-05-top-front-resnet-2cam/status.json @@ -0,0 +1,51 @@ +{ + "suite_name": "2026-04-05-top-front-resnet-2cam", + "updated_at": "2026-04-05 08:52:12", + "phase": "running", + "baseline_reference": { + "source_run": "imf-p1-ph16-ex08-emb384-l12-ms50k-5880g1-20260404-131223", + "best_rollout_avg_reward": 610.8, + "best_step": 21874, + "notes": "Same IMF baseline as Phase-1 best, but switch cameras from [r_vis, top, front] to [top, front] and keep the original ResNet vision backbone." + }, + "smoke_test": { + "status": "passed", + "run_dir": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/smoke-topfront-resnet-ph16-ex08-20260405-085000", + "batch_size": 80, + "num_workers": 4, + "max_steps": 2, + "note": "2-step local CUDA smoke passed without OOM using top/front only." + }, + "main_run": { + "status": "running", + "host": "local", + "gpu": 0, + "pid": 1693348, + "run_name": "imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023", + "run_dir": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023", + "log_path": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/runs/imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023/train_vla.log", + "launch_log": "/home/droid/project/roboimi/.worktrees/feat-imf-attnres-policy/experiment_suites/2026-04-05-top-front-resnet-2cam/launch_logs/imf-resnet-topfront-2cam-ph16-ex08-emb384-l12-ms50k-5090-20260405-085023.launch.log", + "dataset_dir": "/home/droid/project/diana_sim/sim_transfer", + "camera_names": [ + "top", + "front" + ], + "pred_horizon": 16, + "num_action_steps": 8, + "head_n_emb": 384, + "head_n_layer": 12, + "vision_backbone_mode": "resnet", + "pretrained_backbone_weights": null, + "freeze_backbone": false, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/vi77mn5dwd19z4nttxab8", + "latest_step": 500, + "latest_loss": 0.0978, + "process_running": true + } +} diff --git a/experiment_suites/2026-04-05-top-only-resnet-1cam/CHECKLIST.md b/experiment_suites/2026-04-05-top-only-resnet-1cam/CHECKLIST.md new file mode 100644 index 0000000..764a396 --- /dev/null +++ b/experiment_suites/2026-04-05-top-only-resnet-1cam/CHECKLIST.md @@ -0,0 +1,7 @@ +# CHECKLIST + +- [x] Create run contract +- [x] Remote smoke test passes +- [x] Launch 50k main run +- [x] Record pid / log / SwanLab +- [x] Report status back to user diff --git a/experiment_suites/2026-04-05-top-only-resnet-1cam/PLAN.md b/experiment_suites/2026-04-05-top-only-resnet-1cam/PLAN.md new file mode 100644 index 0000000..b7b4b7e --- /dev/null +++ b/experiment_suites/2026-04-05-top-only-resnet-1cam/PLAN.md @@ -0,0 +1,12 @@ +# PLAN + +## Goal +Train a 50k-step IMF baseline with the original ResNet vision backbone using top only as the only image conditioning. + +## Fixed comparison contract +- same hyperparameters as the active top/front run +- cameras: ['top'] +- num_cams=1 +- head.cond_dim=80 +- host: 100.119.99.14 +- gpu: 4 diff --git a/experiment_suites/2026-04-05-top-only-resnet-1cam/notes.md b/experiment_suites/2026-04-05-top-only-resnet-1cam/notes.md new file mode 100644 index 0000000..a7789a7 --- /dev/null +++ b/experiment_suites/2026-04-05-top-only-resnet-1cam/notes.md @@ -0,0 +1,6 @@ +# Notes + +- 2026-04-05 12:58:22: smoke passed for ['top'] on 100.119.99.14 GPU4. +- 2026-04-05 12:59:24: launched main run `imf-resnet-top-1cam-ph16-ex08-emb384-l12-ms50k-l20g4-20260405-125844`. +- 2026-04-05 13:01:20: latest confirmed progress step=400, loss=0.1233. +- SwanLab: https://swanlab.cn/@game-loader/roboimi-vla/runs/egzo29l3z9ftsaunhf025 diff --git a/experiment_suites/2026-04-05-top-only-resnet-1cam/status.json b/experiment_suites/2026-04-05-top-only-resnet-1cam/status.json new file mode 100644 index 0000000..e10e60f --- /dev/null +++ b/experiment_suites/2026-04-05-top-only-resnet-1cam/status.json @@ -0,0 +1,47 @@ +{ + "suite_name": "2026-04-05-top-only-resnet-1cam", + "updated_at": "2026-04-05 13:01:20", + "phase": "running", + "smoke_test": { + "status": "passed", + "host": "100.119.99.14", + "gpu": 4, + "run_dir": "/home/droid/roboimi_suite_20260404/runs/smoke-toponly-resnet-ph16-ex08-20260405-125812", + "batch_size": 80, + "max_steps": 2, + "note": "2-step remote CUDA smoke passed without OOM." + }, + "main_run": { + "status": "running", + "host": "100.119.99.14", + "gpu": 4, + "launch_pid": 164808, + "pid": 164813, + "run_name": "imf-resnet-top-1cam-ph16-ex08-emb384-l12-ms50k-l20g4-20260405-125844", + "run_dir": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-top-1cam-ph16-ex08-emb384-l12-ms50k-l20g4-20260405-125844", + "log_path": "/home/droid/roboimi_suite_20260404/runs/imf-resnet-top-1cam-ph16-ex08-emb384-l12-ms50k-l20g4-20260405-125844/train_vla.log", + "launch_log": "/home/droid/roboimi_suite_20260404/experiment_suite_launch_logs/imf-resnet-top-1cam-ph16-ex08-emb384-l12-ms50k-l20g4-20260405-125844.launch.log", + "dataset_dir": "/home/droid/sim_dataset/sim_transfer", + "camera_names": [ + "top" + ], + "pred_horizon": 16, + "num_action_steps": 8, + "head_cond_dim": 80, + "head_n_emb": 384, + "head_n_layer": 12, + "vision_backbone_mode": "resnet", + "pretrained_backbone_weights": null, + "freeze_backbone": false, + "batch_size": 80, + "lr": 0.00025, + "num_workers": 12, + "max_steps": 50000, + "rollout_val_freq_epochs": 5, + "rollout_num_episodes": 5, + "swanlab_url": "https://swanlab.cn/@game-loader/roboimi-vla/runs/egzo29l3z9ftsaunhf025", + "latest_step": 400, + "latest_loss": 0.1233, + "process_running": true + } +} diff --git a/generate_dataset_videos.py b/generate_dataset_videos.py new file mode 100644 index 0000000..0adae9f --- /dev/null +++ b/generate_dataset_videos.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python3 +""" +将 HDF5 数据集转换为视频，用于可视化检查 + +功能： +1. 将单个 episode 转换为视频 +2. 对比多个 episode 的视频 +3. 放慢播放速度便于观察 +""" +import os +import h5py +import glob +import cv2 +import numpy as np + + +def episode_to_video(episode_file, output_path, camera='top', fps=30, slow_factor=1): + """ + 将单个 episode 转换为视频 + + Args: + episode_file: HDF5 文件路径 + output_path: 输出视频路径 + camera: 要使用的相机名称 + fps: 帧率 + slow_factor: 慢放倍数（1=正常，2=半速） + """ + try: + with h5py.File(episode_file, 'r') as f: + # 读取图像序列 + img_path = f'/observations/images/{camera}' + + if img_path not in f: + print(f" ❌ 相机 {camera} 不存在") + return False + + images = f[img_path][:] # shape: (T, H, W, C) + qpos = f['/observations/qpos'][:] + actions = f['/action'][:] + + total_frames = len(images) + height, width = images.shape[1], images.shape[2] + + # 创建视频写入器 + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + actual_fps = fps // slow_factor + out = cv2.VideoWriter(output_path, fourcc, actual_fps, (width, height)) + + # 逐帧写入 + for i in range(total_frames): + frame = images[i].astype(np.uint8) + + # 在图像上添加信息 + info_text = [ + f"Episode: {os.path.basename(episode_file).replace('.hdf5', '')}", + f"Frame: {i}/{total_frames}", + f"qpos[0:3]: [{qpos[i, 0]:.2f}, {qpos[i, 1]:.2f}, {qpos[i, 2]:.2f}]", + ] + + for j, text in enumerate(info_text): + cv2.putText(frame, text, (10, 30 + j*30), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + + out.write(frame) + + out.release() + print(f" ✅ 保存: {output_path}") + print(f" 帧数: {total_frames}, 尺寸: {width}x{height}, FPS: {actual_fps}") + return True + + except Exception as e: + print(f" ❌ 错误: {e}") + return False + + +def generate_all_videos(camera='top', num_episodes=5, slow_factor=1): + """生成前 N 个 episode 的视频""" + + dataset_dir = "roboimi/demos/dataset/sim_transfer" + episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) + + if len(episode_files) == 0: + print(f"❌ 没有找到数据文件: {dataset_dir}") + return + + # 创建输出目录 + output_dir = '/tmp/dataset_videos' + os.makedirs(output_dir, exist_ok=True) + + print(f"找到 {len(episode_files)} 个 episode 文件") + print(f"将生成前 {min(num_episodes, len(episode_files))} 个 episode 的视频\n") + + # 生成视频 + for i in range(min(num_episodes, len(episode_files))): + ep_file = episode_files[i] + ep_name = os.path.basename(ep_file).replace('.hdf5', '') + output_path = f"{output_dir}/{ep_name}_{camera}.mp4" + + print(f"[{i+1}/{min(num_episodes, len(episode_files))}] {ep_name}") + episode_to_video(ep_file, output_path, camera=camera, slow_factor=slow_factor) + print() + + print(f"✅ 所有视频已保存到: {output_dir}") + print(f"\n播放方法:") + print(f" # 播放单个视频") + print(f" vlc {output_dir}/*.mp4") + print(f" ") + print(f" # 或用文件管理器") + print(f" nautilus {output_dir}") + + +def generate_multi_camera_video(episode_idx=0, slow_factor=1): + """生成包含多个相机的视频（分屏显示）""" + + dataset_dir = "roboimi/demos/dataset/sim_transfer" + episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) + + if episode_idx >= len(episode_files): + print(f"❌ Episode {episode_idx} 不存在") + return + + ep_file = episode_files[episode_idx] + + try: + with h5py.File(ep_file, 'r') as f: + # 获取所有相机 + cameras = [] + for key in f.keys(): + if 'images' in key: + for cam_name in f[key].keys(): + if cam_name not in cameras: + cameras.append(cam_name) + + print(f"Episode {episode_idx} 的相机: {cameras}") + + # 读取所有相机的图像 + all_images = {} + for cam in cameras: + img_path = f'/observations/images/{cam}' + if img_path in f: + all_images[cam] = f[img_path][:] + + if not all_images: + print("❌ 没有找到图像数据") + return + + # 获取第一个相机的尺寸 + first_cam = list(all_images.keys())[0] + total_frames = len(all_images[first_cam]) + height, width = all_images[first_cam].shape[1], all_images[first_cam].shape[2] + + # 创建多相机布局 + num_cams = len(all_images) + cols = min(2, num_cams) + rows = (num_cams + cols - 1) // cols + + canvas_width = width * cols + canvas_height = height * rows + + # 创建视频写入器 + output_path = f'/tmp/dataset_videos/episode_{episode_idx}_all_cameras.mp4' + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, 30 // slow_factor, (canvas_width, canvas_height)) + + # 逐帧合成 + for i in range(total_frames): + canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) + + for cam_idx, cam_name in enumerate(all_images.keys()): + img = all_images[cam_name][i] + + # 计算在画布上的位置 + row = cam_idx // cols + col = cam_idx % cols + y_start = row * height + y_end = y_start + height + x_start = col * width + x_end = x_start + width + + # 调整大小（如果需要） + if img.shape[:2] != (height, width): + img = cv2.resize(img, (width, height)) + + # 放到画布上 + canvas[y_start:y_end, x_start:x_end] = img + + # 添加相机名称 + cv2.putText(canvas, cam_name, (x_start + 10, y_start + 30), + cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2) + + # 添加帧信息 + cv2.putText(canvas, f"Frame: {i}/{total_frames}", (10, canvas_height - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + + out.write(canvas) + + out.release() + print(f"✅ 保存多相机视频: {output_path}") + + except Exception as e: + print(f"❌ 错误: {e}") + + +def compare_episodes(camera='top', slow_factor=2): + """并排对比多个 episode 的视频""" + + dataset_dir = "roboimi/demos/dataset/sim_transfer" + episode_files = sorted(glob.glob(os.path.join(dataset_dir, "episode_*.hdf5"))) + + # 选择要对比的 episode + episodes_to_compare = [0, 1, 2, 3, 4] # 对比前 5 个 + + print(f"对比 Episodes: {episodes_to_compare}") + + # 读取所有 episode 的数据 + all_data = [] + for ep_idx in episodes_to_compare: + if ep_idx >= len(episode_files): + continue + + try: + with h5py.File(episode_files[ep_idx], 'r') as f: + img_path = f'/observations/images/{camera}' + if img_path in f: + all_data.append({ + 'idx': ep_idx, + 'images': f[img_path][:], + 'qpos': f['/observations/qpos'][:] + }) + except: + pass + + if len(all_data) == 0: + print("❌ 没有数据") + return + + # 获取参数 + first_data = all_data[0] + height, width = first_data['images'].shape[1], first_data['images'].shape[2] + total_frames = min([d['images'].shape[0] for d in all_data]) + + # 创建并排布局 + num_compare = len(all_data) + canvas_width = width * num_compare + canvas_height = height + + # 创建视频 + output_path = f'/tmp/dataset_videos/compare_{camera}.mp4' + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, 30 // slow_factor, (canvas_width, canvas_height)) + + print(f"生成对比视频，共 {total_frames} 帧...") + + # 逐帧对比 + for i in range(total_frames): + canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) + + for j, data in enumerate(all_data): + img = data['images'][i] + qpos = data['qpos'][i] + + # 调整大小（如果需要） + if img.shape[:2] != (height, width): + img = cv2.resize(img, (width, height)) + + # 放到画布上 + x_start = j * width + x_end = x_start + width + canvas[:, x_start:x_end] = img + + # 添加信息 + ep_name = f"Ep {data['idx']}" + cv2.putText(canvas, ep_name, (x_start + 10, 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2) + cv2.putText(canvas, f"qpos[0:3]: [{qpos[0]:.2f}, {qpos[1]:.2f}, {qpos[2]:.2f}]", + (x_start + 10, height - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + + # 添加帧号 + cv2.putText(canvas, f"Frame: {i}/{total_frames}", (10, canvas_height - 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) + + out.write(canvas) + + if i % 100 == 0: + print(f" 进度: {i}/{total_frames}") + + out.release() + print(f"✅ 保存对比视频: {output_path}") + + +if __name__ == "__main__": + import sys + + print("="*60) + print("数据集视频生成工具") + print("="*60) + + if len(sys.argv) > 1: + command = sys.argv[1] + + if command == 'compare': + # 对比多个 episode + camera = sys.argv[2] if len(sys.argv) > 2 else 'top' + compare_episodes(camera=camera, slow_factor=2) + + elif command == 'multi': + # 多相机视频 + ep_idx = int(sys.argv[2]) if len(sys.argv) > 2 else 0 + generate_multi_camera_video(episode_idx=ep_idx, slow_factor=1) + + else: + print("未知命令") + else: + # 默认：生成前 5 个 episode 的视频 + print("\n生成前 5 个 episode 的视频（top 相机，慢放 2x）...") + print("="*60 + "\n") + generate_all_videos(camera='top', num_episodes=5, slow_factor=2) + + print("\n" + "="*60) + print("其他用法:") + print(" python generate_dataset_videos.py compare top # 对比多个 episode") + print(" python generate_dataset_videos.py multi 0 # 多相机视频") + print("="*60) diff --git a/roboimi/gr00t/main.py b/gr00t/main.py similarity index 100% rename from roboimi/gr00t/main.py rename to gr00t/main.py diff --git a/roboimi/gr00t/models/__init__.py b/gr00t/models/__init__.py similarity index 100% rename from roboimi/gr00t/models/__init__.py rename to gr00t/models/__init__.py diff --git a/roboimi/detr/models/backbone.py b/gr00t/models/backbone.py similarity index 100% rename from roboimi/detr/models/backbone.py rename to gr00t/models/backbone.py diff --git a/roboimi/gr00t/models/dit.py b/gr00t/models/dit.py similarity index 100% rename from roboimi/gr00t/models/dit.py rename to gr00t/models/dit.py diff --git a/roboimi/gr00t/models/gr00t.py b/gr00t/models/gr00t.py similarity index 100% rename from roboimi/gr00t/models/gr00t.py rename to gr00t/models/gr00t.py diff --git a/roboimi/gr00t/models/modules.py b/gr00t/models/modules.py similarity index 100% rename from roboimi/gr00t/models/modules.py rename to gr00t/models/modules.py diff --git a/roboimi/detr/models/position_encoding.py b/gr00t/models/position_encoding.py similarity index 100% rename from roboimi/detr/models/position_encoding.py rename to gr00t/models/position_encoding.py diff --git a/roboimi/gr00t/policy.py b/gr00t/policy.py similarity index 100% rename from roboimi/gr00t/policy.py rename to gr00t/policy.py diff --git a/roboimi/.gitattributes b/roboimi/.gitattributes new file mode 100644 index 0000000..580d310 --- /dev/null +++ b/roboimi/.gitattributes @@ -0,0 +1 @@ +*.safetensors filter=lfs diff=lfs merge=lfs -text diff --git a/roboimi/__init__.py b/roboimi/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml b/roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml index 1668bc8..7b5e55a 100644 --- a/roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml +++ b/roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml @@ -76,7 +76,7 @@ - + diff --git a/roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml b/roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml new file mode 100644 index 0000000..e532054 --- /dev/null +++ b/roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml @@ -0,0 +1,6 @@ + + + + + + diff --git a/roboimi/assets/models/manipulators/DianaMed/box.xml b/roboimi/assets/models/manipulators/DianaMed/box.xml index f351cc3..c016926 100644 --- a/roboimi/assets/models/manipulators/DianaMed/box.xml +++ b/roboimi/assets/models/manipulators/DianaMed/box.xml @@ -3,7 +3,7 @@ - + diff --git a/roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml b/roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml new file mode 100644 index 0000000..642bd78 --- /dev/null +++ b/roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml @@ -0,0 +1,19 @@ + + + + + + + + + + + + + + + + + + + diff --git a/roboimi/assets/models/manipulators/DianaMed/table_square.xml b/roboimi/assets/models/manipulators/DianaMed/table_square.xml index 4813a53..fb03d1f 100644 --- a/roboimi/assets/models/manipulators/DianaMed/table_square.xml +++ b/roboimi/assets/models/manipulators/DianaMed/table_square.xml @@ -7,6 +7,6 @@ - + diff --git a/roboimi/assets/robots/arm_base.py b/roboimi/assets/robots/arm_base.py index 5cf94bd..0e80f7b 100644 --- a/roboimi/assets/robots/arm_base.py +++ b/roboimi/assets/robots/arm_base.py @@ -1,8 +1,46 @@ import mujoco import numpy as np +from pathlib import Path from roboimi.utils.KDL_utils import KDL_utils +def resolve_robot_asset_path(asset_path): + if asset_path is None: + return None + + raw_path = Path(asset_path).expanduser() + if raw_path.is_absolute(): + return str(raw_path.resolve()) + + current_dir = Path(__file__).resolve().parent + package_root = current_dir.parents[1] + repo_root = current_dir.parents[2] + + candidates = [] + if raw_path.parts and raw_path.parts[0] == 'roboimi': + candidates.append(repo_root / raw_path) + + candidates.extend([ + current_dir / raw_path, + package_root / raw_path, + repo_root / raw_path, + ]) + + normalized_candidates = [] + seen = set() + for candidate in candidates: + resolved = candidate.resolve() + if resolved not in seen: + normalized_candidates.append(resolved) + seen.add(resolved) + + for candidate in normalized_candidates: + if candidate.exists(): + return str(candidate) + + return str(normalized_candidates[0]) + + class ArmBase(object): def __init__(self, name=None, @@ -11,8 +49,8 @@ class ArmBase(object): gripper=None ): self.name = name - self.urdf_path = urdf_path - self.xml_path = xml_path + self.urdf_path = resolve_robot_asset_path(urdf_path) + self.xml_path = resolve_robot_asset_path(xml_path) self.gripper = gripper self.robot_model = mujoco.MjModel.from_xml_path(filename=self.xml_path, assets=None) self.robot_data = mujoco.MjData(self.robot_model) diff --git a/roboimi/assets/robots/diana_med.py b/roboimi/assets/robots/diana_med.py index 234b50e..691837e 100644 --- a/roboimi/assets/robots/diana_med.py +++ b/roboimi/assets/robots/diana_med.py @@ -58,8 +58,8 @@ class BiDianaMed(ArmBase): def __init__(self): super().__init__( name="Bidiana", - urdf_path="./assets/models/manipulators/DianaMed/DualDianaMed.urdf", - xml_path="./assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml", + urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf", + xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml", gripper=None ) self.left_arm = self.Arm(self, 'single', self.urdf_path) @@ -90,4 +90,40 @@ class BiDianaMed(ArmBase): def init_qpos(self): """ Robot's init joint position. """ return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0]) + + +class BiDianaMedSocketPeg(ArmBase): + def __init__(self): + super().__init__( + name="Bidiana_socket_peg", + urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf", + xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml", + gripper=None + ) + self.left_arm = self.Arm(self, 'single', self.urdf_path) + self.left_arm.set_Arm_base_link('left_base_link') + self.left_arm.set_Arm_ee_link('left_link7') + self.left_arm.InitKDL + self.left_arm.joint_index = ['l_j1','l_j2','l_j3','l_j4','l_j5','l_j6','l_j7'] + self.left_arm.gripper_index = ['l_finger_joint_left','r_finger_joint_left'] + self.left_arm.actuator_index = ['a1_l','a2_l','a3_l','a4_l','a5_l','a6_l','a7_l','gripper_left'] + self.left_arm.setArmInitPose(self.init_qpos) + self.arms.append(self.left_arm) + self.right_arm = self.Arm(self,'single', self.urdf_path) + self.right_arm.set_Arm_base_link('right_base_link') + self.right_arm.set_Arm_ee_link('right_link7') + self.right_arm.InitKDL + self.right_arm.joint_index = ['r_j1','r_j2','r_j3','r_j4','r_j5','r_j6','r_j7'] + self.right_arm.gripper_index = ['l_finger_joint_right','r_finger_joint_right'] + self.right_arm.actuator_index = ['a1_r','a2_r','a3_r','a4_r','a5_r','a6_r','a7_r','gripper_right'] + self.right_arm.setArmInitPose(self.init_qpos) + self.arms.append(self.right_arm) + self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num + self.kp = 500 * np.ones(self.jnt_num) + self.kd = 44.57 * np.ones(self.jnt_num) + + @property + def init_qpos(self): + """ Robot's init joint position. """ + return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0]) diff --git a/roboimi/demos/config.yaml b/roboimi/demos/config.yaml index 3b16eb1..cf754d7 100644 --- a/roboimi/demos/config.yaml +++ b/roboimi/demos/config.yaml @@ -38,8 +38,13 @@ episode_len: # leave empty here by default camera_names: [] # leave empty here by default xml_dir: # leave empty here by default +# action smoothing settings (for GR00T) +use_action_smoothing: true +smooth_method: "ema" # Options: "ema", "moving_avg", "lowpass", "none" +smooth_alpha: 0.3 # Smoothing factor (0-1), smaller = smoother + # transformer settings -batch_size: 15 +batch_size: 10 state_dim: 16 action_dim: 16 lr_backbone: 0.00001 diff --git a/roboimi/demos/diana_air_insert_policy.py b/roboimi/demos/diana_air_insert_policy.py new file mode 100644 index 0000000..f8ffaa4 --- /dev/null +++ b/roboimi/demos/diana_air_insert_policy.py @@ -0,0 +1,203 @@ +import numpy as np +from pyquaternion import Quaternion + +from roboimi.demos.diana_policy import PolicyBase + + +class TestAirInsertPolicy(PolicyBase): + ACTION_OBJECT_Z_OFFSET = 0.078 + SOCKET_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64) + PEG_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64) + SOCKET_OUTER_GRASP_STRATEGY = "socket_outer" + LEGACY_GRASP_STRATEGY = "legacy" + SOCKET_HOLD_Z = 0.85 + PEG_INSERT_START_OFFSET = np.array([0.105, 0.0, 0.0], dtype=np.float64) + INSERT_START_T = 650 + INSERT_END_T = 730 + LEFT_SOCKET_GRIPPER_CLOSED = -100 + RIGHT_PEG_GRIPPER_CLOSED = -100 + SOCKET_APPROACH_Z = 1.05 + EPISODE_END_T = 1000 + + def __init__(self, inject_noise=False, grasp_strategy=SOCKET_OUTER_GRASP_STRATEGY): + super().__init__(inject_noise=inject_noise) + valid_strategies = { + self.SOCKET_OUTER_GRASP_STRATEGY, + self.LEGACY_GRASP_STRATEGY, + } + if grasp_strategy not in valid_strategies: + raise ValueError( + f"Unsupported air insert grasp_strategy={grasp_strategy!r}; " + f"expected one of {sorted(valid_strategies)}" + ) + self.grasp_strategy = grasp_strategy + + def generate_trajectory(self, task_state): + return self._generate_socket_peg_trajectory(task_state) + + def _generate_socket_peg_trajectory(self, task_state): + socket_xyz = np.asarray(task_state["socket_pos"], dtype=np.float64) + peg_xyz = np.asarray(task_state["peg_pos"], dtype=np.float64) + + init_mocap_pose_left = np.array( + [ + -0.17297014, + 1.00485877, + 1.32773627, + 7.06825181e-01, + 8.20281078e-06, + -7.07388269e-01, + -5.20399313e-06, + ], + dtype=np.float64, + ) + init_mocap_pose_right = np.array( + [ + 0.17297014, + 0.9951369, + 1.32773623, + 2.59463975e-06, + 7.07388269e-01, + 5.59551158e-06, + 7.06825181e-01, + ], + dtype=np.float64, + ) + + left_init_quat = Quaternion(init_mocap_pose_left[3:]) + right_init_quat = Quaternion(init_mocap_pose_right[3:]) + + left_pick_quat = ( + left_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45) + ).elements + right_pick_quat = ( + right_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45) + ).elements + + socket_hold_action = np.array( + [socket_xyz[0] - 0.078, socket_xyz[1], self.SOCKET_HOLD_Z], dtype=np.float64 + ) + + peg_init_xyz = peg_xyz + np.array( + [0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET + 0.01] + ) + peg_lift_center = np.array( + [peg_xyz[0] + 0.078, socket_hold_action[1], self.SOCKET_HOLD_Z - 0.01], + dtype=np.float64, + ) + # The front camera looks along +Y, so visual right-to-left insertion is + # world +X -> -X. With the socket XML in identity orientation, its + # tunnel axis is local/world X, so the peg approaches from +X and stops + # when its leading face reaches the socket's internal pin. + peg_insert_end_center = np.array( + [ + socket_hold_action[0] + 0.078 * 2 + 0.04 + 0.06 - 0.01, + socket_hold_action[1], + self.SOCKET_HOLD_Z - 0.01, + ], + dtype=np.float64, + ) + + self.left_trajectory = [ + { + "t": 1, + "xyz": init_mocap_pose_left[:3], + "quat": init_mocap_pose_left[3:], + "gripper": 100, + }, + { + "t": 130, + "xyz": socket_xyz + + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]), + "quat": left_pick_quat, + "gripper": 100, + }, + { + "t": 180, + "xyz": socket_xyz + + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]), + "quat": left_pick_quat, + "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED, + }, + { + "t": 450, + "xyz": socket_hold_action, + "quat": left_pick_quat, + "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED, + }, + { + "t": 750, + "xyz": socket_hold_action, + "quat": left_pick_quat, + "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED, + }, + { + "t": self.EPISODE_END_T, + "xyz": socket_hold_action, + "quat": left_pick_quat, + "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED, + }, + ] + + self.right_trajectory = [ + { + "t": 1, + "xyz": init_mocap_pose_right[:3], + "quat": init_mocap_pose_right[3:], + "gripper": 100, + }, + { + "t": 80, + "xyz": peg_init_xyz, + "quat": right_pick_quat, + "gripper": 100, + }, + { + "t": 150, + "xyz": peg_init_xyz, + "quat": right_pick_quat, + "gripper": 100, + }, + { + "t": 180, + "xyz": peg_init_xyz, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + { + "t": 450, + "xyz": peg_init_xyz, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + { + "t": 550, + "xyz": peg_lift_center, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + { + "t": self.INSERT_START_T, + "xyz": peg_lift_center, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + { + "t": self.INSERT_END_T, + "xyz": peg_insert_end_center, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + { + "t": 750, + "xyz": peg_insert_end_center, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + { + "t": self.EPISODE_END_T, + "xyz": peg_insert_end_center, + "quat": right_pick_quat, + "gripper": self.RIGHT_PEG_GRIPPER_CLOSED, + }, + ] diff --git a/roboimi/demos/diana_eval.py b/roboimi/demos/diana_eval.py deleted file mode 100644 index a5e71e5..0000000 --- a/roboimi/demos/diana_eval.py +++ /dev/null @@ -1,118 +0,0 @@ -import torch -import os -import numpy as np -import matplotlib.pyplot as plt -from tqdm import tqdm -from einops import rearrange -from roboimi.utils.utils import set_seed -from roboimi.utils.io_utils import IOUtils -from roboimi.utils.model_interface import ModelInterface -from roboimi.envs.double_pos_ctrl_env import make_sim_env -# from visualize_episodes import save_videos -from roboimi.utils.act_ex_utils import sample_transfer_pose - - - -#should be added into IOUtils -def get_image(obs,camera_names): - curr_images = [] - for cam_name in camera_names: - curr_image = rearrange(obs['images'][cam_name], 'h w c -> c h w') - curr_images.append(curr_image) - curr_image = np.stack(curr_images, axis=0) - curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0) - return curr_image - - -def eval_bc(config, ckpt_name='policy_best.ckpt', save_episode=True): - set_seed(1) - model_interface = ModelInterface(config) - model_interface.setup() - policy = IOUtils.load_policy(config, ckpt_name) - stats = IOUtils.load_stats(config['ckpt_dir']) - num_rollouts = 3 - episode_returns = [] - highest_rewards = [] - - - - - - run_episode(config, policy, stats, - save_episode,num_rollouts) - # episode_return, episode_highest_reward = run_episode(config, policy, stats, - # save_episode,num_rollouts) - - - - -def run_episode(config, policy, stats, save_episode,num_rollouts): - - if 'sim_transfer' in config['task_name']: - task_name = 'sim_transfer' #config['task_name'] - env = make_sim_env(task_name) - - max_timesteps = config['episode_len'] - max_timesteps = int(max_timesteps * 1) - pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std'] - post_process = lambda a: a * stats['action_std'] + stats['action_mean'] - box_pos = sample_transfer_pose() - for rollout_id in range(num_rollouts): - print("\nrollout_id===",rollout_id,"\n") - image_list = [] - rewards = [] - query_frequency = config['policy_config'].get('num_queries', 1) - print("query_freq =====",query_frequency) - env.reset(box_pos) - with torch.inference_mode(): - for t in range(700): - image_list.append(env._get_image_obs()['images'] if 'images' in env._get_image_obs() else {print("img error")}) - qpos_numpy = np.array(env._get_qpos_obs()['qpos']) - qpos = pre_process(qpos_numpy) - qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0) - curr_image = get_image(env._get_image_obs(), config['camera_names']) - if config['policy_class'] in ["ACT", "ACTTV", "GR00T"]: - if t % query_frequency == 0: - all_actions = policy(qpos, curr_image) - raw_action = all_actions[:, t % query_frequency] - raw_action = raw_action.squeeze(0).cpu().numpy() - elif config['policy_class'] == "CNNMLP": - raw_action = policy(qpos, curr_image) - else: - raise NotImplementedError - - - action = post_process(raw_action) - print("action == ",action) - env.step_jnt(action) - rewards.append(env.rew) - env.render() - - - rewards = np.array(rewards) - # episode_return = np.sum(rewards[rewards != None]) - # episode_highest_reward = np.max(rewards) - # env.viewer.close() - - # del env - # return episode_return, episode_highest_reward - - - - -def test_env(): - try: - env = make_sim_env('sim_transfer') - env.reset() - while True: pass - except KeyboardInterrupt: - del env - print("stop") - -if __name__ == '__main__': - # test_env() - io_utils = IOUtils() - config = io_utils.load_config() - eval_bc(config) - - diff --git a/roboimi/demos/diana_policy.py b/roboimi/demos/diana_policy.py index f710d76..7c847c5 100644 --- a/roboimi/demos/diana_policy.py +++ b/roboimi/demos/diana_policy.py @@ -104,8 +104,8 @@ class TestPickAndTransferPolicy(PolicyBase): {"t": 1, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": -100}, # sleep {"t": 75, "xyz": np.array([(0.8+box_xyz[0])*0.5,(1.0+box_xyz[1])*0.5,init_mocap_pose_right[2]]), "quat": gripper_approach_quat.elements, "gripper": 100}, {"t": 225, "xyz": box_xyz + np.array([0, 0, 0.3]), "quat": gripper_pick_quat.elements, "gripper": 100}, # approach the cube - {"t": 275, "xyz": box_xyz + np.array([0, 0, 0.12]), "quat": gripper_pick_quat.elements, "gripper": 100}, # go down - {"t": 280, "xyz": box_xyz + np.array([0, 0, 0.12]), "quat": gripper_pick_quat.elements, "gripper": -100}, # close gripper + {"t": 275, "xyz": box_xyz + np.array([0, 0, 0.11]), "quat": gripper_pick_quat.elements, "gripper": 100}, # go down + {"t": 280, "xyz": box_xyz + np.array([0, 0, 0.11]), "quat": gripper_pick_quat.elements, "gripper": -100}, # close gripper {"t": 450, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": -100},# approach wait position {"t": 500, "xyz": meet_xyz + np.array([0.1, 0, 0.0]), "quat": meet_right_quat.elements, "gripper": -100},# approach meet position {"t": 510, "xyz": meet_xyz + np.array([0.1, 0, 0.0]), "quat": meet_right_quat.elements, "gripper": 100}, # open gripper @@ -116,8 +116,8 @@ class TestPickAndTransferPolicy(PolicyBase): self.left_trajectory = [ {"t": 1, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": -100},# sleep {"t": 250, "xyz": meet_xyz + np.array([-0.5, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": 100}, # approach meet position - {"t": 500, "xyz": meet_xyz + np.array([-0.15, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": 100}, # move to meet position - {"t": 505, "xyz": meet_xyz + np.array([-0.15, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # close gripper + {"t": 500, "xyz": meet_xyz + np.array([-0.14, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": 100}, # move to meet position + {"t": 505, "xyz": meet_xyz + np.array([-0.14, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # close gripper {"t": 675, "xyz": meet_xyz + np.array([-0.3, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # move left {"t": 700, "xyz": meet_xyz + np.array([-0.3, 0, 0.0]), "quat": meet_left_quat.elements, "gripper": -100}, # stay ] diff --git a/roboimi/demos/diana_record_sim_episodes.py b/roboimi/demos/diana_record_sim_episodes.py index 5eadf79..1b0dad3 100644 --- a/roboimi/demos/diana_record_sim_episodes.py +++ b/roboimi/demos/diana_record_sim_episodes.py @@ -1,29 +1,47 @@ import time -import os,collections,sys +import os import numpy as np -import h5py from roboimi.envs.double_pos_ctrl_env import make_sim_env -from diana_policy import TestPickAndTransferPolicy +from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy +from roboimi.demos.diana_policy import TestPickAndTransferPolicy import cv2 -from roboimi.utils.act_ex_utils import sample_transfer_pose +from roboimi.utils.act_ex_utils import sample_air_insert_socket_peg_state, sample_transfer_pose +from roboimi.utils.constants import SIM_TASK_CONFIGS +from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter import pathlib HOME_PATH = str(pathlib.Path(__file__).parent.resolve()) DATASET_DIR = HOME_PATH + '/dataset' -def main(): - task_name = 'sim_transfer' - dataset_dir = DATASET_DIR + '/sim_transfer' #SIM_TASK_CONFIGS[task_name]['dataset_dir'] - num_episodes = 100 #SIM_TASK_CONFIGS[task_name]['num_episodes'] - onscreen_render = None #config['onscreen_render'] - inject_noise = False - render_cam_name = 'angle' - - episode_len = 700 #SIM_TASK_CONFIGS[task_name]['episode_len'] - camera_names = ['angle','r_vis', 'top'] #SIM_TASK_CONFIGS[task_name]['camera_names'] +def sample_task_state(task_name): if task_name == 'sim_transfer': - policy = TestPickAndTransferPolicy(inject_noise) + return sample_transfer_pose() + if task_name == 'sim_air_insert_socket_peg': + return sample_air_insert_socket_peg_state() + raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}') + + +def make_policy(task_name, inject_noise=False, grasp_strategy=None): + if task_name == 'sim_transfer': + return TestPickAndTransferPolicy(inject_noise) + if task_name == 'sim_air_insert_socket_peg': + if grasp_strategy is None: + return TestAirInsertPolicy(inject_noise) + return TestAirInsertPolicy(inject_noise, grasp_strategy=grasp_strategy) + raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}') + + +def main(task_name='sim_transfer'): + task_cfg = SIM_TASK_CONFIGS[task_name] + dataset_dir = task_cfg['dataset_dir'] + num_episodes = 100 + inject_noise = False + + episode_len = task_cfg['episode_len'] + camera_names = task_cfg['camera_names'] + image_size = (256, 256) + if task_name in {'sim_transfer', 'sim_air_insert_socket_peg'}: print(task_name) else: raise NotImplementedError @@ -31,64 +49,46 @@ def main(): success = [] env = make_sim_env(task_name) - policy = TestPickAndTransferPolicy(inject_noise) + policy = make_policy(task_name, inject_noise=inject_noise) + + # 等待osmesa完全启动后再开始收集数据 + print("等待osmesa线程启动...") + time.sleep(60) + print("osmesa已就绪，开始收集数据...") + for episode_idx in range(num_episodes): - obs = [] - reward_ee = [] + sum_reward = 0.0 + max_reward = float('-inf') print(f'\n{episode_idx=}') print('Rollout out EE space scripted policy') - box_pos = sample_transfer_pose() - env.reset(box_pos) + task_state = sample_task_state(task_name) + env.reset(task_state) + episode_writer = StreamingEpisodeWriter( + dataset_path=os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5'), + max_timesteps=episode_len, + camera_names=camera_names, + image_size=image_size, + ) for step in range(episode_len): - - - action = policy.predict(box_pos,step) - env.step(action) + raw_action = policy.predict(task_state, step) + env.step(raw_action) env.render() - reward_ee.append(env.rew) - obs.append(env.obs) - sum_reward = np.sum(reward_ee) - max_reward = np.max(reward_ee) + sum_reward += env.rew + max_reward = max(max_reward, env.rew) + episode_writer.append( + qpos=env.obs['qpos'], + action=raw_action, + images=env.obs['images'], + ) if max_reward == env.max_reward: success.append(1) print(f"{episode_idx=} Successful, {sum_reward=}") - t0 = time.time() - data_dict = { - '/observations/qpos': [], - '/action': [], - } - - for cam_name in camera_names: - data_dict[f'/observations/images/{cam_name}'] = [] - for i in range(episode_len): - print("type qpos==",obs[i]['qpos']) - data_dict['/observations/qpos'].append(obs[i]['qpos']) - data_dict['/action'].append(obs[i]['action']) - for cam_name in camera_names: - data_dict[f'/observations/images/{cam_name}'].append(obs[i]['images'][cam_name]) - - dataset_path = os.path.join(dataset_dir, f'episode_{episode_idx}') - - with h5py.File(dataset_path + '.hdf5', 'w', rdcc_nbytes=1024 ** 2 * 2) as root: - max_timesteps = episode_len - root.attrs['sim'] = True - obs_ = root.create_group('observations') - image = obs_.create_group('images') - for cam_name in camera_names: - _ = image.create_dataset(cam_name, (max_timesteps, 480, 640, 3), dtype='uint8', - chunks=(1, 480, 640, 3), ) - qpos = obs_.create_dataset('qpos', (max_timesteps, 16)) - action = root.create_dataset('action', (max_timesteps, 16)) - for name, array in data_dict.items(): - root[name][...] = np.array(array) + episode_writer.commit() else: success.append(0) print(f"{episode_idx=} Failed") print(max_reward) - del obs - del reward_ee - del sum_reward - del max_reward + episode_writer.discard() # del policy # env.viewer.close() @@ -102,4 +102,4 @@ def main(): if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/roboimi/demos/eval.py b/roboimi/demos/eval.py deleted file mode 100644 index 792c81a..0000000 --- a/roboimi/demos/eval.py +++ /dev/null @@ -1,152 +0,0 @@ -import torch -import os -import numpy as np -import matplotlib.pyplot as plt -from tqdm import tqdm -from einops import rearrange -from roboimi.utils.utils import set_seed -from roboimi.utils.io_utils import IOUtils -from roboimi.utils.model_interface import ModelInterface -from roboimi.envs.vx300s_jnt import make_sim_env -import time - -# from visualize_episodes import save_videos -from roboimi.utils.utils import sample_box_pose, sample_insertion_pose - - - -#should be added into IOUtils -def get_image(obs,camera_names): - curr_images = [] - for cam_name in camera_names: - curr_image = rearrange(obs['images'][cam_name], 'h w c -> c h w') - curr_images.append(curr_image) - curr_image = np.stack(curr_images, axis=0) - curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0) - return curr_image - - -def eval_bc(config, ckpt_name='policy_best.ckpt', save_episode=True): - set_seed(1) - model_interface = ModelInterface(config) - task_name = 'sim_insertion' #config['task_name'] - model_interface.setup() - policy = IOUtils.load_policy(config, ckpt_name) - stats = IOUtils.load_stats(config['ckpt_dir']) - num_rollouts = 3 - episode_returns = [] - highest_rewards = [] - for rollout_id in range(num_rollouts): - episode_return, episode_highest_reward = run_episode(config, policy, stats, - save_episode,rollout_id) - - - - -def run_episode(config, policy, stats, save_episode,rollout_id): - print("\nrollout_id===",rollout_id,"\n") - pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std'] - post_process = lambda a: a * stats['action_std'] + stats['action_mean'] - if 'sim_insertion' in config['task_name']: - peg_pose, socket_pose = sample_insertion_pose() - box_pose = np.hstack((peg_pose[:3],socket_pose[:3])) # used in sim reset - task_name = 'sim_insertion' #config['task_name'] - env = make_sim_env(task_name) - env.reset(box_pose) - max_timesteps = config['episode_len'] - max_timesteps = int(max_timesteps * 1) - - image_list = [] - rewards = [] - query_frequency = config['policy_config'].get('num_queries', 1) - - with torch.inference_mode(): - for t in range(700): - # print("obs_img",env.obs['images']) - image_list.append(env.obs['images'] if 'images' in env.obs else {print("img error")}) - qpos_numpy = np.array(env.obs['qpos']) - qpos = pre_process(qpos_numpy) - qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0) - curr_image = get_image(env.obs, config['camera_names']) - if config['policy_class'] == "ACT" or "ACTTV": - if t % query_frequency == 0: - all_actions = policy(qpos, curr_image) - elif config['policy_class'] == "CNNMLP": - raw_action = policy(qpos, curr_image) - else: - raise NotImplementedError - raw_action = all_actions[:, t % query_frequency] - raw_action = raw_action.squeeze(0).cpu().numpy() - action = post_process(raw_action) - - env.step(action) - rewards.append(env.rew) - env.render() - - - rewards = np.array(rewards) - episode_return = np.sum(rewards[rewards != None]) - episode_highest_reward = np.max(rewards) - env.viewer.close() - - del env - return episode_return, episode_highest_reward - - -def test_env(): - try: - env = make_sim_env('sim_insertion') - box_pos = np.concatenate(sample_insertion_pose()) - env.reset(box_pos) - while True: pass - except KeyboardInterrupt: - del env - print("stop") - - -if __name__ == '__main__': - test_env() - # io_utils = IOUtils() - # config = io_utils.load_config() - # eval_bc(config) - - - - -# config===== {'onscreen_render': False, -# 'eval': 1, -# 'ckpt_dir': 'ckpt_models', -# 'num_epochs': 3000, -# 'temporal_agg': False, -# 'policy_class': 'ACT', -# 'backbone': 'resnet18', -# 'seed': 0, 'real_robot': 0, -# 'task_name': 'sim_insertion', -# 'images_render_height': 480, -# 'images_render_width': 640, -# 'left_arm_DOF_number': 6, -# 'right_arm_DOF_number': 6, -# 'left_qpos_raw': 8, -# 'right_qpos_raw': 8, -# 'left_qvel_raw': 8, -# 'right_qvel_raw': 8, -# 'dataset_dir': '/home/arm/lzd/act_env/dataset/sim_insertion', -# 'num_episodes': 7, -# 'episode_len': 400, -# 'camera_names': ['top'], -# 'xml_dir': None, -# 'batch_size': 8, -# 'state_dim': 14, -# 'action_dim': 14, -# 'lr_backbone': 1e-05, -# 'enc_layers': 4, -# 'dec_layers': 7, -# 'nheads': 8, -# 'qpos_noise_std': 0, -# 'DT': 0.02, -# 'lr': 1e-05, -# 'kl_weight': 10, -# 'chunk_size': 100, -# 'hidden_dim': 512, -# 'dim_feedforward': 3200, -# 'policy_config': {'lr': 1e-05, 'num_queries': 100, 'kl_weight': 10, 'hidden_dim': 512, 'dim_feedforward': 3200, 'lr_backbone': 1e-05, 'backbone': 'resnet18', 'enc_layers': 4, 'dec_layers': 7, 'nheads': 8, 'camera_names': ['top']}} \ No newline at end of file diff --git a/roboimi/demos/training.py b/roboimi/demos/training.py deleted file mode 100644 index 858960b..0000000 --- a/roboimi/demos/training.py +++ /dev/null @@ -1,179 +0,0 @@ -import torch -import os -from tqdm import tqdm -import numpy as np -from copy import deepcopy -from itertools import repeat -import matplotlib.pyplot as plt -import time -from roboimi.utils.utils import set_seed, compute_dict_mean, detach_dict, load_data -from roboimi.utils.io_utils import IOUtils -from roboimi.utils.model_interface import ModelInterface -import matplotlib.pyplot as plt - -def train_bc(config): - num_epochs = config['num_epochs'] - ckpt_dir = config['ckpt_dir'] - seed = config['seed'] - - os.makedirs(ckpt_dir, exist_ok=True) - - set_seed(seed) - - model_interface = ModelInterface(config) - model_interface.setup() - - policy = model_interface.make_policy() - policy.cuda() - optimizer = model_interface.make_optimizer(policy) - # print("cam names=====",config['camera_names']) - train_dataloader, val_dataloader, stats, _ = load_data( - config['dataset_dir'], - config['num_episodes'], - config['camera_names'], - config['batch_size'], - config['batch_size']) - - IOUtils.save_stats(ckpt_dir, stats) - - train_history = [] - validation_history = [] - min_val_loss = np.inf - min_train_loss = np.inf - best_ckpt_info = None - - plt.ion() - fig, ax = plt.subplots() - train_losses, val_losses = [], [] - train_line, = ax.plot([], [], label='Train Loss') - val_line, = ax.plot([], [], label='Validation Loss') - ax.autoscale_view() - ax.set_xlabel('Epoch') - ax.set_ylabel('Loss') - ax.legend() - ax.grid(True) - - - train_annotation = ax.annotate('', xy=(0, 0), textcoords='offset points') - val_annotation = ax.annotate('', xy=(0, 0), textcoords='offset points') - - - min_train_text = ax.text(0.85, 0.5, '', transform=ax.transAxes, fontsize=10, verticalalignment='center', horizontalalignment='left', bbox=dict(facecolor='white', alpha=0.5)) - min_val_text = ax.text(0.85, 0.45, '', transform=ax.transAxes, fontsize=10, verticalalignment='center', horizontalalignment='left', bbox=dict(facecolor='white', alpha=0.5)) - - for epoch in tqdm(range(num_epochs)): - print(f'\nEpoch {epoch}') - - # Validation - epoch_val_loss, epoch_summary = validate(policy, val_dataloader) - validation_history.append(epoch_summary) - val_losses.append(epoch_val_loss.cpu().item()) - - if epoch_val_loss < min_val_loss: - min_val_loss = epoch_val_loss - min_val_epoch = epoch - best_ckpt_info = (epoch, min_val_loss, - deepcopy(policy.state_dict())) - - print(f'Val loss: {epoch_val_loss:.5f}') - print_summary(epoch_summary) - - # Training - epoch_train_loss, epoch_summary = train_epoch( - policy, optimizer, train_dataloader) - train_history.append(epoch_summary) - train_losses.append(epoch_train_loss.cpu().item()) - - if epoch_train_loss < min_train_loss: - min_train_loss = epoch_train_loss - min_train_epoch = epoch - - print(f'Train loss: {epoch_train_loss:.5f}') - print_summary(epoch_summary) - - # Update the plot with the new data - train_line.set_xdata(range(len(train_losses))) - train_line.set_ydata(train_losses) - val_line.set_xdata(range(len(val_losses))) - val_line.set_ydata(val_losses) - - # Update annotations with the latest loss values at their respective positions - train_annotation.set_position((len(train_losses)-1, train_losses[-1])) - train_annotation.xy = (len(train_losses)-1, train_losses[-1]) - train_annotation.set_text(f'{train_losses[-1]:.5f}') - - val_annotation.set_position((len(val_losses)-1, val_losses[-1])) - val_annotation.xy = (len(val_losses)-1, val_losses[-1]) - val_annotation.set_text(f'{val_losses[-1]:.5f}') - - # Update text objects with the minimum loss values, fixed on the right side - min_train_text.set_text(f'Min Train Loss: {min_train_loss:.5f} (Epoch {min_train_epoch})') - min_val_text.set_text(f'Min Val Loss: {min_val_loss:.5f} (Epoch {min_val_epoch})') - - ax.relim() - ax.autoscale_view() - plt.draw() - plt.pause(0.1) - - - plt.ioff() - IOUtils.save_checkpoint(policy, 'last', ckpt_dir, seed, 'last') - - best_epoch, min_val_loss, best_state_dict = best_ckpt_info - IOUtils.save_checkpoint(best_state_dict, best_epoch, - ckpt_dir, seed, 'best', min_val_loss) - print( - f'Training finished:\nSeed {seed}, val loss {min_val_loss:.6f} at epoch {best_epoch}') - - IOUtils.plot_history(train_history, validation_history, - num_epochs, ckpt_dir, seed) - - return best_ckpt_info - - - - - - -def validate(policy, dataloader): - policy.eval() - epoch_dicts = [] - with torch.inference_mode(): - for data in dataloader: - forward_dict = forward_pass(data, policy) - epoch_dicts.append(forward_dict) - epoch_summary = compute_dict_mean(epoch_dicts) - return epoch_summary['loss'], epoch_summary - - -def train_epoch(policy, optimizer, dataloader): - policy.train() - epoch_dicts = [] - for data in dataloader: - optimizer.zero_grad() - forward_dict = forward_pass(data, policy) - loss = forward_dict['loss'] - loss.backward() - optimizer.step() - epoch_dicts.append(detach_dict(forward_dict)) - epoch_summary = compute_dict_mean(epoch_dicts) - return epoch_summary['loss'], epoch_summary - - -def forward_pass(data, policy): - image_data, qpos_data, action_data, is_pad = data - image_data, qpos_data, action_data, is_pad = image_data.cuda( - ), qpos_data.cuda(), action_data.cuda(), is_pad.cuda() - return policy(qpos_data, image_data, action_data, is_pad) - - -def print_summary(summary): - summary_string = ' '.join( - [f'{k}: {v.item():.3f}' for k, v in summary.items()]) - print(summary_string) - - -if __name__ == '__main__': - io_utils = IOUtils() - config = io_utils.load_config() - train_bc(config) diff --git a/roboimi/demos/view_raw_action_trajectory.py b/roboimi/demos/view_raw_action_trajectory.py new file mode 100644 index 0000000..f44d756 --- /dev/null +++ b/roboimi/demos/view_raw_action_trajectory.py @@ -0,0 +1,36 @@ +import argparse +import numpy as np + +from roboimi.utils.raw_action_trajectory_viewer import launch_raw_action_trajectory_viewer + + +def parse_args(): + parser = argparse.ArgumentParser(description="Launch an interactive MuJoCo viewer with raw-action trajectory overlay.") + parser.add_argument("trajectory_path", help="Path to raw_action.npy or trajectory.npz") + parser.add_argument("--task-name", default="sim_transfer") + parser.add_argument("--line-radius", type=float, default=0.004) + parser.add_argument("--max-markers", type=int, default=1500) + parser.add_argument( + "--box-pos", + type=float, + nargs=3, + default=None, + help="Optional box xyz to use when resetting the environment", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + box_pos = np.asarray(args.box_pos, dtype=np.float32) if args.box_pos is not None else None + launch_raw_action_trajectory_viewer( + args.trajectory_path, + task_name=args.task_name, + line_radius=args.line_radius, + max_markers=args.max_markers, + box_pos=box_pos, + ) + + +if __name__ == "__main__": + main() diff --git a/roboimi/demos/vla_scripts/eval_vla.py b/roboimi/demos/vla_scripts/eval_vla.py new file mode 100644 index 0000000..437ea7d --- /dev/null +++ b/roboimi/demos/vla_scripts/eval_vla.py @@ -0,0 +1,968 @@ +""" +VLA 策略评估脚本（简化版） + +该脚本使用 agent 内置的队列管理来评估训练好的 VLA 策略。 +无需单独的评估器类 - agent 处理一切！ + +使用方法: + python roboimi/demos/eval_vla_simple.py + python roboimi/demos/eval_vla_simple.py eval.ckpt_path=checkpoints/vla_model_final.pt + python roboimi/demos/eval_vla_simple.py eval.ckpt_path=checkpoints/vla_model_best.pt +""" + +import sys +import os +import json +import logging +import time +import torch +import numpy as np +import hydra +from pathlib import Path +from typing import Any, Dict, Optional +from tqdm import tqdm +from omegaconf import DictConfig, OmegaConf +from hydra.utils import instantiate +from einops import rearrange + +from roboimi.envs.double_pos_ctrl_env import make_sim_env +from roboimi.utils.act_ex_utils import ( + sample_air_insert_socket_peg_state, + sample_transfer_pose, +) +from roboimi.vla.eval_utils import execute_policy_action + +sys.path.append(os.getcwd()) + +log = logging.getLogger(__name__) + +if not OmegaConf.has_resolver("len"): + OmegaConf.register_new_resolver("len", lambda x: len(x)) + + +def _configure_headless_mujoco_gl(eval_cfg: DictConfig) -> None: + if not bool(eval_cfg.get('headless', False)): + return + if os.environ.get('MUJOCO_GL'): + return + os.environ['MUJOCO_GL'] = 'egl' + log.info('headless eval detected; set MUJOCO_GL=egl') + + +def make_sim_env(task_name: str, headless: bool = False): + from roboimi.envs.double_pos_ctrl_env import make_sim_env as _make_sim_env_impl + return _make_sim_env_impl(task_name, headless=headless) + + +def load_checkpoint( + ckpt_path: str, + agent_cfg: DictConfig, + device: str = 'cuda' +) -> torch.nn.Module: + """ + 从检查点加载训练好的 VLA 模型，使用 Hydra agent 配置。 + + Args: + ckpt_path: 检查点文件路径 (.pt) + agent_cfg: Hydra agent 配置，用于实例化 + device: 加载模型的设备 + + Returns: + 加载后的 VLAAgent 模型 + """ + from pathlib import Path as PathLib + + ckpt_path = PathLib(ckpt_path).absolute() + if not ckpt_path.exists(): + raise FileNotFoundError(f"检查点未找到: {ckpt_path}") + + log.info(f"从 {ckpt_path} 加载检查点") + checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False) + log.info(f"检查点键值: {checkpoint.keys()}") + + # 加载数据集统计信息用于归一化 + stats = checkpoint.get('dataset_stats', None) + + # 使用数据集统计信息从 Hydra 配置实例化 agent + log.info("从配置实例化 agent...") + agent = instantiate(agent_cfg, dataset_stats=stats) + + # 加载模型状态 + agent.load_state_dict(checkpoint['model_state_dict']) + log.info(f"✅ 模型状态已加载 (步数: {checkpoint.get('step', 'unknown')})") + + if stats is not None: + log.info(f"✅ 数据集统计信息已加载 (归一化: {stats.get('normalization_type', 'gaussian')})") + else: + # 后备方案：尝试从外部 JSON 文件加载（兼容旧检查点） + stats_path = ckpt_path.parent / 'dataset_stats.json' + if stats_path.exists(): + with open(stats_path, 'r') as f: + stats = json.load(f) + log.info("✅ 数据集统计信息已从外部 JSON 加载（旧版本兼容）") + else: + log.warning("⚠️ 未找到数据集统计信息。动作将无法反归一化！") + + agent.eval() + agent.to(device) + + log.info(f"✅ 模型已成功加载到 {device}") + return agent, stats + + +def prepare_observation( + obs: Dict, + camera_names: list, + image_resize_shape: Optional[tuple[int, int]] = (224, 224), +) -> Dict: + """ + 将环境观测转换为 agent 格式。 + + Args: + obs: 环境观测字典，包含图像和 qpos + camera_names: 摄像头名称列表 + + Returns: + agent 格式的观测字典 + """ + # 转换图像: numpy -> tensor, HWC -> CHW + images = {} + for cam_name in camera_names: + img = obs['images'][cam_name] + if image_resize_shape is not None: + import cv2 + img = cv2.resize(img, tuple(image_resize_shape), interpolation=cv2.INTER_LINEAR) + img = rearrange(img, 'h w c -> c h w') + img = torch.from_numpy(img / 255.0).float() + images[cam_name] = img + + # 转换 qpos: numpy -> tensor + qpos = torch.from_numpy(obs['qpos']).float() + + return {'qpos': qpos, 'images': images} + + +def _to_numpy_action(action: Any) -> np.ndarray: + if isinstance(action, torch.Tensor): + return action.detach().cpu().numpy().astype(np.float32, copy=True) + return np.asarray(action, dtype=np.float32).copy() + + +def _mean_or_zero(values: list[float]) -> float: + return float(np.mean(values)) if values else 0.0 + + +def _stats_or_zero(values: list[float]) -> dict[str, float]: + if not values: + return { + 'mean': 0.0, + 'std': 0.0, + 'min': 0.0, + 'max': 0.0, + } + array = np.asarray(values, dtype=np.float64) + return { + 'mean': float(array.mean()), + 'std': float(array.std()), + 'min': float(array.min()), + 'max': float(array.max()), + } + + +def _summarize_timing_breakdown( + all_timings: dict[str, list[float]], + model_forward_flags: list[bool], +) -> dict[str, Any]: + model_forward_flags = [bool(flag) for flag in model_forward_flags] + return { + 'count': int(len(model_forward_flags)), + 'model_forward_count': int(sum(model_forward_flags)), + 'all_steps_ms': { + stage: _stats_or_zero(values) + for stage, values in all_timings.items() + }, + 'model_forward_steps_ms': { + stage: _stats_or_zero( + [value for value, should_keep in zip(values, model_forward_flags) if should_keep] + ) + for stage, values in all_timings.items() + }, + } + + +def _json_friendly(value: Any) -> Any: + if isinstance(value, dict): + return {str(key): _json_friendly(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_friendly(item) for item in value] + if isinstance(value, Path): + return str(value) + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, (np.integer, np.floating)): + return value.item() + return value + + +def _resolve_artifact_paths(eval_cfg: DictConfig) -> dict[str, Optional[str]]: + save_timing = bool(eval_cfg.get('save_timing', False)) + save_trajectory = bool( + eval_cfg.get('save_trajectory', False) or eval_cfg.get('save_trajectory_npz', False) + ) + save_trajectory_image = bool(eval_cfg.get('save_trajectory_image', False)) + wants_artifacts = any([ + bool(eval_cfg.get('save_artifacts', False)), + save_timing, + save_trajectory, + save_trajectory_image, + bool(eval_cfg.get('record_video', False)), + ]) + output_dir: Optional[Path] = None + if wants_artifacts: + artifact_dir = eval_cfg.get('artifact_dir', None) + if artifact_dir: + output_dir = Path(str(artifact_dir)).expanduser().resolve() + else: + ckpt_stem = Path(str(eval_cfg.ckpt_path)).stem or 'rollout' + timestamp = time.strftime('%Y%m%d-%H%M%S') + output_dir = (Path.cwd() / 'rollout_artifacts' / f'{ckpt_stem}-{timestamp}').resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + video_camera_name = None + if bool(eval_cfg.get('record_video', False)): + configured_camera_name = eval_cfg.get('video_camera_name', None) + if configured_camera_name is None: + configured_camera_name = eval_cfg.get('video_camera', None) + if configured_camera_name is not None: + video_camera_name = str(configured_camera_name) + elif eval_cfg.get('camera_names'): + video_camera_name = str(eval_cfg.camera_names[0]) + else: + raise ValueError('record_video=true requires eval.video_camera_name or a non-empty eval.camera_names') + + trajectory_image_camera_name = None + if save_trajectory_image: + configured_camera_name = eval_cfg.get('trajectory_image_camera_name', None) + if configured_camera_name is None: + configured_camera_name = eval_cfg.get('trajectory_image_camera', None) + if configured_camera_name is not None: + trajectory_image_camera_name = str(configured_camera_name) + elif eval_cfg.get('camera_names'): + camera_names = [str(name) for name in eval_cfg.camera_names] + trajectory_image_camera_name = 'front' if 'front' in camera_names else camera_names[0] + else: + raise ValueError( + 'save_trajectory_image=true requires eval.trajectory_image_camera_name ' + 'or a non-empty eval.camera_names' + ) + + return { + 'output_dir': str(output_dir) if output_dir is not None else None, + 'summary_json': ( + str(output_dir / 'rollout_summary.json') + if output_dir is not None and bool(eval_cfg.get('save_summary_json', False)) + else None + ), + 'timing_json': ( + str(output_dir / 'timing.json') + if output_dir is not None and save_timing + else None + ), + 'trajectory_npz': ( + str(output_dir / 'trajectory.npz') + if output_dir is not None and save_trajectory + else None + ), + 'video_mp4': ( + str(output_dir / f'rollout_{video_camera_name}.mp4') + if output_dir is not None and bool(eval_cfg.get('record_video', False)) + and video_camera_name is not None + else None + ), + 'video_camera_name': video_camera_name, + 'trajectory_image_camera_name': trajectory_image_camera_name, + } + + +def _get_video_frame(obs: Dict, camera_name: Optional[str]) -> Optional[np.ndarray]: + if camera_name is None: + return None + frame = obs['images'][camera_name] + frame = np.asarray(frame) + if frame.ndim != 3 or frame.shape[2] != 3: + raise ValueError( + f'Video frame for camera {camera_name} must have shape (H, W, 3), got {frame.shape}' + ) + if frame.dtype != np.uint8: + frame = np.clip(frame, 0, 255).astype(np.uint8) + return frame + + +def _open_video_writer(output_path: str, frame_size: tuple[int, int], fps: int): + import cv2 + + output_path = str(output_path) + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + writer = cv2.VideoWriter(output_path, fourcc, float(fps), frame_size) + if not writer.isOpened(): + raise RuntimeError(f'无法打开视频输出: {output_path}') + return writer + + +def _episode_trajectory_image_path( + artifact_paths: dict[str, Optional[str]], + episode_idx: int, +) -> Optional[str]: + output_dir = artifact_paths.get('output_dir') + camera_name = artifact_paths.get('trajectory_image_camera_name') + if output_dir is None or camera_name is None: + return None + return str(Path(output_dir) / f'rollout_{camera_name}_ep{episode_idx + 1:02d}_trajectory.png') + + +def _build_action_trajectory_positions(raw_actions: list[np.ndarray]) -> dict[str, np.ndarray]: + if not raw_actions: + empty = np.zeros((0, 3), dtype=np.float32) + return {'left': empty, 'right': empty} + raw_action_array = np.asarray(raw_actions, dtype=np.float32) + return { + 'left': raw_action_array[:, :3].astype(np.float32, copy=True), + 'right': raw_action_array[:, 7:10].astype(np.float32, copy=True), + } + + +def _append_capsule_markers_to_scene(scene, markers: list[dict]) -> None: + import mujoco + + for marker in markers: + if scene.ngeom >= scene.maxgeom: + break + geom = scene.geoms[scene.ngeom] + mujoco.mjv_initGeom( + geom, + mujoco.mjtGeom.mjGEOM_CAPSULE, + np.zeros(3, dtype=np.float64), + np.zeros(3, dtype=np.float64), + np.eye(3, dtype=np.float64).reshape(-1), + np.asarray(marker['rgba'], dtype=np.float32), + ) + mujoco.mjv_connector( + geom, + mujoco.mjtGeom.mjGEOM_CAPSULE, + float(marker['radius']), + np.asarray(marker['from'], dtype=np.float64), + np.asarray(marker['to'], dtype=np.float64), + ) + scene.ngeom += 1 + + +def _save_rollout_trajectory_image( + env, + output_path: Optional[str], + raw_actions: list[np.ndarray], + camera_name: Optional[str], + *, + line_radius: float = 0.004, + max_markers: int = 1500, +) -> Optional[str]: + if output_path is None or camera_name is None: + return None + + # IMPORTANT: + # Keep this import lazy so headless rollout can set MUJOCO_GL=egl before + # anything imports mujoco. Importing this helper at module import time would + # pull in mujoco too early on remote headless hosts and make rollout fail + # with gladLoadGL / missing DISPLAY errors. + from roboimi.utils.raw_action_trajectory_viewer import build_trajectory_capsule_markers + + output_path = str(output_path) + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + + frame = None + owned_renderer = None + positions = _build_action_trajectory_positions(raw_actions) + markers = build_trajectory_capsule_markers( + positions, + max_markers=max_markers, + radius=line_radius, + ) + + try: + renderer = None + if callable(getattr(env, '_get_or_create_offscreen_renderer', None)): + renderer = env._get_or_create_offscreen_renderer() + elif hasattr(env, 'mj_model') and hasattr(env, 'mj_data'): + import mujoco + + renderer = mujoco.Renderer(env.mj_model, height=480, width=640) + owned_renderer = renderer + + if renderer is not None and hasattr(env, 'mj_data'): + renderer.update_scene(env.mj_data, camera=str(camera_name)) + if markers: + _append_capsule_markers_to_scene(renderer.scene, markers) + frame = renderer.render()[:, :, ::-1] + finally: + if owned_renderer is not None: + owned_renderer.close() + + if frame is None and callable(getattr(env, '_get_image_obs', None)): + obs = env._get_image_obs() + frame = _get_video_frame(obs, str(camera_name)) + + if frame is None: + return None + + import cv2 + + cv2.imwrite(output_path, frame) + return output_path + + +class _RolloutVideoRecorder: + def __init__(self, output_path: Optional[str], fps: int): + self.output_path = output_path + self.fps = int(fps) + self.writer = None + + def write(self, frame: Optional[np.ndarray]): + if self.output_path is None or frame is None: + return + if self.writer is None: + frame_size = (int(frame.shape[1]), int(frame.shape[0])) + self.writer = _open_video_writer(self.output_path, frame_size, self.fps) + self.writer.write(frame) + + def close(self): + if self.writer is not None: + self.writer.release() + self.writer = None + + +def _read_body_pose(env, body_name: str): + try: + if callable(getattr(env, 'getBodyPos', None)) and callable(getattr(env, 'getBodyQuat', None)): + pos = env.getBodyPos(body_name) + quat = env.getBodyQuat(body_name) + else: + body = env.mj_data.body(body_name) + pos = body.xpos + quat = body.xquat + except Exception: + return None + + return { + 'pos': np.asarray(pos, dtype=np.float32).copy(), + 'quat': np.asarray(quat, dtype=np.float32).copy(), + } + + +def _get_executed_ee_poses(env) -> dict[str, np.ndarray]: + candidates = { + 'left_link7': ('left_link7', 'eef_left'), + 'right_link7': ('right_link7', 'eef_right'), + 'eef_left': ('eef_left', 'left_link7'), + 'eef_right': ('eef_right', 'right_link7'), + } + poses = {} + for body_key, body_names in candidates.items(): + pose = None + for body_name in body_names: + pose = _read_body_pose(env, body_name) + if pose is not None: + break + if pose is None: + pose = { + 'pos': np.full(3, np.nan, dtype=np.float32), + 'quat': np.full(4, np.nan, dtype=np.float32), + } + poses[f'{body_key}_pos'] = pose['pos'] + poses[f'{body_key}_quat'] = pose['quat'] + return poses + + +def _empty_rollout_trajectory() -> dict[str, list]: + return { + 'episode_index': [], + 'step': [], + 'reward': [], + 'raw_action': [], + 'applied_action': [], + 'executed_left_link7_pos': [], + 'executed_left_link7_quat': [], + 'executed_right_link7_pos': [], + 'executed_right_link7_quat': [], + 'executed_eef_left_pos': [], + 'executed_eef_left_quat': [], + 'executed_eef_right_pos': [], + 'executed_eef_right_quat': [], + 'model_inference_triggered': [], + 'obs_read_time_ms': [], + 'preprocess_time_ms': [], + 'inference_time_ms': [], + 'env_step_time_ms': [], + 'total_time_ms': [], + } + + +def _append_rollout_step( + storage: dict[str, list], + episode_index: int, + timestep: int, + reward: Optional[float], + raw_action: np.ndarray, + executed_action: np.ndarray, + executed_poses: dict[str, np.ndarray], + timing_ms: dict[str, float], + model_inference_triggered: bool, +): + storage['episode_index'].append(int(episode_index)) + storage['step'].append(int(timestep)) + storage['reward'].append(float(reward) if reward is not None else np.nan) + storage['raw_action'].append(raw_action.astype(np.float32, copy=True)) + storage['applied_action'].append(executed_action.astype(np.float32, copy=True)) + storage['executed_left_link7_pos'].append(executed_poses['left_link7_pos']) + storage['executed_left_link7_quat'].append(executed_poses['left_link7_quat']) + storage['executed_right_link7_pos'].append(executed_poses['right_link7_pos']) + storage['executed_right_link7_quat'].append(executed_poses['right_link7_quat']) + storage['executed_eef_left_pos'].append(executed_poses['eef_left_pos']) + storage['executed_eef_left_quat'].append(executed_poses['eef_left_quat']) + storage['executed_eef_right_pos'].append(executed_poses['eef_right_pos']) + storage['executed_eef_right_quat'].append(executed_poses['eef_right_quat']) + storage['model_inference_triggered'].append(bool(model_inference_triggered)) + for key, value in timing_ms.items(): + storage[key].append(float(value)) + + +def _save_rollout_trajectory_npz(output_path: str, storage: dict[str, list]): + step = np.asarray(storage['step'], dtype=np.int32) + raw_action = np.asarray(storage['raw_action'], dtype=np.float32) + applied_action = np.asarray(storage['applied_action'], dtype=np.float32) + executed_left_link7_pos = np.asarray(storage['executed_left_link7_pos'], dtype=np.float32) + executed_left_link7_quat = np.asarray(storage['executed_left_link7_quat'], dtype=np.float32) + executed_right_link7_pos = np.asarray(storage['executed_right_link7_pos'], dtype=np.float32) + executed_right_link7_quat = np.asarray(storage['executed_right_link7_quat'], dtype=np.float32) + executed_eef_left_pos = np.asarray(storage['executed_eef_left_pos'], dtype=np.float32) + executed_eef_left_quat = np.asarray(storage['executed_eef_left_quat'], dtype=np.float32) + executed_eef_right_pos = np.asarray(storage['executed_eef_right_pos'], dtype=np.float32) + executed_eef_right_quat = np.asarray(storage['executed_eef_right_quat'], dtype=np.float32) + np.savez_compressed( + output_path, + episode_index=np.asarray(storage['episode_index'], dtype=np.int32), + step=step, + timestep=step, + reward=np.asarray(storage['reward'], dtype=np.float32), + raw_action=raw_action, + raw_predicted_ee_action=raw_action, + applied_action=applied_action, + executed_ee_action=applied_action, + executed_left_link7_pos=executed_left_link7_pos, + executed_left_link7_quat=executed_left_link7_quat, + executed_right_link7_pos=executed_right_link7_pos, + executed_right_link7_quat=executed_right_link7_quat, + executed_eef_left_pos=executed_eef_left_pos, + executed_eef_left_quat=executed_eef_left_quat, + executed_eef_right_pos=executed_eef_right_pos, + executed_eef_right_quat=executed_eef_right_quat, + left_ee_pos=executed_eef_left_pos, + left_ee_quat=executed_eef_left_quat, + right_ee_pos=executed_eef_right_pos, + right_ee_quat=executed_eef_right_quat, + model_inference_triggered=np.asarray(storage['model_inference_triggered'], dtype=bool), + obs_read_time_ms=np.asarray(storage['obs_read_time_ms'], dtype=np.float32), + preprocess_time_ms=np.asarray(storage['preprocess_time_ms'], dtype=np.float32), + inference_time_ms=np.asarray(storage['inference_time_ms'], dtype=np.float32), + env_step_time_ms=np.asarray(storage['env_step_time_ms'], dtype=np.float32), + total_time_ms=np.asarray(storage['total_time_ms'], dtype=np.float32), + ) + + +def _save_summary_json(output_path: str, summary: dict[str, Any]): + with open(output_path, 'w', encoding='utf-8') as f: + json.dump(_json_friendly(summary), f, ensure_ascii=False, indent=2) + + +class ActionSmoother: + """ + 动作平滑器（指数移动平均） + 用于平滑执行动作以获得更稳定的控制 + """ + + def __init__(self, alpha: float = 0.3): + """ + Args: + alpha: 平滑系数 (0-1)，值越大越重视当前动作 + """ + self.alpha = alpha + self.prev_action = None + + def smooth(self, action: np.ndarray) -> np.ndarray: + """ + 平滑动作 + + Args: + action: 当前动作 + + Returns: + 平滑后的动作 + """ + if self.prev_action is None: + smoothed = action + else: + smoothed = self.alpha * action + (1 - self.alpha) * self.prev_action + self.prev_action = smoothed + return smoothed + + def reset(self): + """重置平滑器状态""" + self.prev_action = None + + +def _close_env(env): + if env is None: + return + + if hasattr(env, 'exit_flag'): + env.exit_flag = True + + cam_thread = getattr(env, 'cam_thread', None) + if cam_thread is not None and hasattr(cam_thread, 'join'): + cam_thread.join(timeout=1.0) + + viewer = getattr(env, 'viewer', None) + if viewer is not None and hasattr(viewer, 'close'): + viewer.close() + + +def _sample_task_reset_state(task_name: str): + if task_name == 'sim_air_insert_socket_peg': + return sample_air_insert_socket_peg_state() + if 'sim_transfer' in task_name: + return sample_transfer_pose() + raise NotImplementedError(f'Unsupported eval task reset sampling: {task_name}') + + +def _run_eval(cfg: DictConfig): + """ + 使用 agent 内置队列管理的简化版 VLA 评估 + + 所有评估参数来自 vla/conf/eval.yaml，合并到 cfg 中。 + 命令行覆盖: python eval_vla_simple.py eval.ckpt_path=... eval.num_episodes=5 + """ + + # 打印配置 + print("=" * 80) + print("VLA 评估配置:") + print("=" * 80) + print(OmegaConf.to_yaml(cfg)) + print("=" * 80) + + eval_cfg = cfg.eval + _configure_headless_mujoco_gl(eval_cfg) + device = eval_cfg.device + camera_names = list(eval_cfg.camera_names) + artifact_paths = _resolve_artifact_paths(eval_cfg) + video_recorder = _RolloutVideoRecorder( + output_path=artifact_paths['video_mp4'], + fps=int(eval_cfg.get('video_fps', 30)), + ) + rollout_trajectory = _empty_rollout_trajectory() + global_obs_read_times_ms = [] + global_preprocess_times_ms = [] + global_inference_times_ms = [] + global_env_step_times_ms = [] + global_total_times_ms = [] + global_model_forward_flags = [] + + # ========================================================================= + # 加载模型 + # ========================================================================= + log.info(f"🚀 从 {eval_cfg.ckpt_path} 加载模型...") + agent, dataset_stats = load_checkpoint( + ckpt_path=eval_cfg.ckpt_path, + agent_cfg=cfg.agent, + device=device + ) + vision_encoder = getattr(agent, 'vision_encoder', None) + image_resize_shape = getattr(vision_encoder, 'eval_image_resize_shape', (224, 224)) + + # 重置 agent 的队列 + agent.reset() + + # 可选：动作平滑器 + smoother = ActionSmoother(alpha=eval_cfg.smooth_alpha) if eval_cfg.use_smoothing else None + + # ========================================================================= + # 创建环境 + # ========================================================================= + env = make_sim_env(eval_cfg.task_name, headless=eval_cfg.headless) + + # ========================================================================= + # 运行评估回合 + # ========================================================================= + all_stats = [] + episode_rewards = [] + episode_max_rewards = [] + try: + for episode_idx in range(eval_cfg.num_episodes): + print(f"\n{'='*60}") + print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}") + print(f"{'='*60}\n") + + task_state = _sample_task_reset_state(str(eval_cfg.task_name)) + env.reset(task_state) + + # 为新回合重置 agent 队列 + agent.reset() + if smoother: + smoother.reset() + + # 计时统计 + obs_read_times_ms = [] + preprocess_times_ms = [] + inference_times_ms = [] + env_step_times_ms = [] + total_times_ms = [] + model_forward_flags = [] + episode_reward = 0.0 + episode_max_reward = float('-inf') + episode_raw_actions: list[np.ndarray] = [] + + with torch.inference_mode(): + for t in tqdm(range(eval_cfg.max_timesteps), desc=f"回合 {episode_idx + 1}"): + start_total = time.perf_counter() + + # 从环境获取观测 + obs = env._get_image_obs() + qpos_obs = env._get_qpos_obs() + obs['qpos'] = qpos_obs['qpos'] + end_obs_read = time.perf_counter() + + video_frame = _get_video_frame(obs, artifact_paths['video_camera_name']) + video_recorder.write(video_frame) + + # 准备给 agent 的观测 + observation = prepare_observation( + obs, + camera_names, + image_resize_shape=image_resize_shape, + ) + end_preprocess = time.perf_counter() + + # 选择动作（agent 内部处理队列管理） + action_queue = getattr(agent, '_queues', {}).get('action', None) + model_inference_triggered = len(action_queue) == 0 if action_queue is not None else True + start_inference = time.perf_counter() + action = agent.select_action(observation) + + if str(device).startswith('cuda') and torch.cuda.is_available(): + torch.cuda.synchronize() + end_inference = time.perf_counter() + + # 转换为 numpy + raw_action = _to_numpy_action(action) + episode_raw_actions.append(raw_action.astype(np.float32, copy=True)) + + # 调试：打印当前时间步的动作（由配置控制） + if eval_cfg.get('verbose_action', False): + print(f"\n[Step {t:3d}] 预测动作: {raw_action}") + print(f" - 动作形状: {raw_action.shape}") + print(f" - 动作范围: [{raw_action.min():.4f}, {raw_action.max():.4f}]") + print(f" - 动作均值: {raw_action.mean():.4f}, 标准差: {raw_action.std():.4f}") + + # 可选：平滑动作 + executed_action = raw_action.copy() + if smoother: + executed_action = smoother.smooth(executed_action) + + # 执行动作 + start_env_step = time.perf_counter() + execute_policy_action(env, executed_action) + end_env_step = time.perf_counter() + executed_poses = _get_executed_ee_poses(env) + reward = getattr(env, 'rew', None) + if reward is not None: + reward = float(reward) + episode_reward += reward + episode_max_reward = max(episode_max_reward, reward) + if not eval_cfg.headless: + env.render() + + end_total = time.perf_counter() + + step_timing_ms = { + 'obs_read_time_ms': (end_obs_read - start_total) * 1000.0, + 'preprocess_time_ms': (end_preprocess - end_obs_read) * 1000.0, + 'inference_time_ms': (end_inference - start_inference) * 1000.0, + 'env_step_time_ms': (end_env_step - start_env_step) * 1000.0, + 'total_time_ms': (end_total - start_total) * 1000.0, + } + + # 记录计时 + obs_read_times_ms.append(step_timing_ms['obs_read_time_ms']) + preprocess_times_ms.append(step_timing_ms['preprocess_time_ms']) + inference_times_ms.append(step_timing_ms['inference_time_ms']) + env_step_times_ms.append(step_timing_ms['env_step_time_ms']) + total_times_ms.append(step_timing_ms['total_time_ms']) + model_forward_flags.append(bool(model_inference_triggered)) + global_obs_read_times_ms.append(step_timing_ms['obs_read_time_ms']) + global_preprocess_times_ms.append(step_timing_ms['preprocess_time_ms']) + global_inference_times_ms.append(step_timing_ms['inference_time_ms']) + global_env_step_times_ms.append(step_timing_ms['env_step_time_ms']) + global_total_times_ms.append(step_timing_ms['total_time_ms']) + global_model_forward_flags.append(bool(model_inference_triggered)) + + if artifact_paths['trajectory_npz'] is not None: + _append_rollout_step( + rollout_trajectory, + episode_index=episode_idx, + timestep=t, + reward=reward, + raw_action=raw_action, + executed_action=executed_action, + executed_poses=executed_poses, + timing_ms=step_timing_ms, + model_inference_triggered=model_inference_triggered, + ) + + # ========================================================================= + # 打印回合统计 + # ========================================================================= + avg_obs_read_time_ms = _mean_or_zero(obs_read_times_ms) + avg_preprocess_time_ms = _mean_or_zero(preprocess_times_ms) + avg_inference_time_ms = _mean_or_zero(inference_times_ms) + avg_env_step_time_ms = _mean_or_zero(env_step_times_ms) + avg_total_time_ms = _mean_or_zero(total_times_ms) + timing_breakdown = _summarize_timing_breakdown( + { + 'obs_read': obs_read_times_ms, + 'preprocess': preprocess_times_ms, + 'inference': inference_times_ms, + 'env_step': env_step_times_ms, + 'loop_total': total_times_ms, + }, + model_forward_flags, + ) + episode_artifact_paths = { + 'video': artifact_paths['video_mp4'], + 'trajectory': artifact_paths['trajectory_npz'], + 'trajectory_image': _save_rollout_trajectory_image( + env, + _episode_trajectory_image_path(artifact_paths, episode_idx), + episode_raw_actions, + artifact_paths['trajectory_image_camera_name'], + ), + 'timing': artifact_paths['timing_json'] or artifact_paths['summary_json'], + } + + stats = { + 'inference_fps': 1000.0 / avg_inference_time_ms if avg_inference_time_ms > 0 else 0.0, + 'control_fps': 1000.0 / avg_total_time_ms if avg_total_time_ms > 0 else 0.0, + 'avg_obs_read_time_ms': avg_obs_read_time_ms, + 'avg_preprocess_time_ms': avg_preprocess_time_ms, + 'avg_inference_time_ms': avg_inference_time_ms, + 'avg_env_step_time_ms': avg_env_step_time_ms, + 'avg_total_time_ms': avg_total_time_ms, + 'num_inferences': int(sum(model_forward_flags)), + 'num_model_forwards': int(sum(model_forward_flags)), + 'num_steps': len(total_times_ms), + 'episode_reward': float(episode_reward), + 'episode_max_reward': ( + float(episode_max_reward) if episode_max_reward != float('-inf') else None + ), + 'artifact_paths': episode_artifact_paths, + 'timing_breakdown_ms': timing_breakdown['all_steps_ms'], + 'timing_summary': timing_breakdown, + } + all_stats.append(stats) + episode_rewards.append(float(episode_reward)) + if episode_max_reward != float('-inf'): + episode_max_rewards.append(float(episode_max_reward)) + + print(f"\n回合 {episode_idx + 1} 完成 ({eval_cfg.max_timesteps} 时间步)") + print(f" 模型推理 FPS: {stats['inference_fps']:.2f} Hz") + print(f" 控制循环 FPS: {stats['control_fps']:.2f} Hz") + print(f" 平均读观测时间: {stats['avg_obs_read_time_ms']:.2f} ms") + print(f" 平均预处理时间: {stats['avg_preprocess_time_ms']:.2f} ms") + print(f" 平均推理时间: {stats['avg_inference_time_ms']:.2f} ms") + print(f" 平均环境步进时间: {stats['avg_env_step_time_ms']:.2f} ms") + print(f" 平均总时间: {stats['avg_total_time_ms']:.2f} ms") + print(f" 总推理次数: {stats['num_inferences']}") + print(f" 回合累计奖励: {stats['episode_reward']:.2f}") + + # ========================================================================= + # 总体统计 + # ========================================================================= + print(f"\n{'='*60}") + print("评估完成!") + print(f"{'='*60}") + + summary = { + 'num_episodes': int(eval_cfg.num_episodes), + 'episode_rewards': episode_rewards, + 'episode_max_rewards': episode_max_rewards, + 'avg_reward': float(np.mean(episode_rewards)) if episode_rewards else 0.0, + 'avg_max_reward': float(np.mean(episode_max_rewards)) if episode_max_rewards else 0.0, + 'episodes': all_stats, + 'artifact_dir': artifact_paths['output_dir'], + 'artifacts': artifact_paths, + } + + if all_stats: + avg_inference_fps = np.mean([s['inference_fps'] for s in all_stats]) + avg_control_fps = np.mean([s['control_fps'] for s in all_stats]) + avg_obs_read_time = _mean_or_zero(global_obs_read_times_ms) + avg_preprocess_time = _mean_or_zero(global_preprocess_times_ms) + avg_inference_time = _mean_or_zero(global_inference_times_ms) + avg_env_step_time = _mean_or_zero(global_env_step_times_ms) + avg_total_time = _mean_or_zero(global_total_times_ms) + summary.update({ + 'avg_inference_fps': float(avg_inference_fps), + 'avg_control_fps': float(avg_control_fps), + 'avg_obs_read_time_ms': float(avg_obs_read_time), + 'avg_preprocess_time_ms': float(avg_preprocess_time), + 'avg_inference_time_ms': float(avg_inference_time), + 'avg_env_step_time_ms': float(avg_env_step_time), + 'avg_total_time_ms': float(avg_total_time), + 'timing_summary': _summarize_timing_breakdown( + { + 'obs_read': global_obs_read_times_ms, + 'preprocess': global_preprocess_times_ms, + 'inference': global_inference_times_ms, + 'env_step': global_env_step_times_ms, + 'loop_total': global_total_times_ms, + }, + global_model_forward_flags, + ), + }) + + print(f"\n总体统计 ({eval_cfg.num_episodes} 个回合):") + print(f" 平均模型推理 FPS: {avg_inference_fps:.2f} Hz") + print(f" 平均控制循环 FPS: {avg_control_fps:.2f} Hz") + print(f" 平均读观测时间: {avg_obs_read_time:.2f} ms") + print(f" 平均预处理时间: {avg_preprocess_time:.2f} ms") + print(f" 平均推理时间: {avg_inference_time:.2f} ms") + print(f" 平均环境步进时间: {avg_env_step_time:.2f} ms") + print(f" 平均总时间: {avg_total_time:.2f} ms") + print(f" 平均累计奖励: {summary['avg_reward']:.2f}") + + if artifact_paths['trajectory_npz'] is not None: + _save_rollout_trajectory_npz(artifact_paths['trajectory_npz'], rollout_trajectory) + if artifact_paths['summary_json'] is not None: + _save_summary_json(artifact_paths['summary_json'], summary) + if artifact_paths['timing_json'] is not None: + _save_summary_json(artifact_paths['timing_json'], summary.get('timing_summary', {})) + print() + return _json_friendly(summary) + finally: + video_recorder.close() + _close_env(env) + + +@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config") +def main(cfg: DictConfig): + return _run_eval(cfg) + + +if __name__ == '__main__': + main() diff --git a/roboimi/demos/vla_scripts/train_vla.py b/roboimi/demos/vla_scripts/train_vla.py new file mode 100644 index 0000000..e4a063c --- /dev/null +++ b/roboimi/demos/vla_scripts/train_vla.py @@ -0,0 +1,986 @@ +import sys +import os +import logging +import json +import pickle +import importlib +import hydra +import torch +import re +from tqdm import tqdm +from omegaconf import DictConfig, OmegaConf +from torch.utils.data import DataLoader, random_split +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from pathlib import Path + +# 确保正确的导入路径（不能依赖 cwd，因为 Hydra 会在运行时切换 cwd） +def _ensure_repo_root_on_syspath(): + repo_root = Path(__file__).resolve().parents[3] + repo_root_str = str(repo_root) + if repo_root_str in sys.path: + sys.path.remove(repo_root_str) + sys.path.insert(0, repo_root_str) + return repo_root + + +_PROBLEMATIC_LD_PRELOAD_SUBSTRINGS = ('/usr/NX/lib/libnxegl.so', 'libnxegl.so') + + +def _clean_ld_preload_value(value: str | None): + if not value: + return value, False + entries = [entry for entry in value.split() if entry] + filtered = [ + entry for entry in entries + if not any(marker in entry for marker in _PROBLEMATIC_LD_PRELOAD_SUBSTRINGS) + ] + changed = filtered != entries + cleaned = ' '.join(filtered) if filtered else None + return cleaned, changed + + +def _maybe_reexec_without_problematic_ld_preload(): + if __name__ != '__main__': + return False + if os.environ.get('_ROBOIMI_LD_PRELOAD_SANITIZED') == '1': + return False + + cleaned, changed = _clean_ld_preload_value(os.environ.get('LD_PRELOAD')) + if not changed: + return False + + new_env = dict(os.environ) + new_env['_ROBOIMI_LD_PRELOAD_SANITIZED'] = '1' + if cleaned: + new_env['LD_PRELOAD'] = cleaned + else: + new_env.pop('LD_PRELOAD', None) + + print( + 'Detected problematic LD_PRELOAD entry for CUDA/cuDNN; re-executing train_vla.py without it.', + file=sys.stderr, + flush=True, + ) + os.execvpe(sys.executable, [sys.executable, *sys.argv], new_env) + + +_REPO_ROOT = _ensure_repo_root_on_syspath() + +from hydra.utils import instantiate + +log = logging.getLogger(__name__) + +# 注册列表长度解析器（用于配置中如 ${len:${data.camera_names}}） +if not OmegaConf.has_resolver("len"): + OmegaConf.register_new_resolver("len", lambda x: len(x)) + + +def _resolve_run_output_dir() -> Path: + try: + from hydra.core.hydra_config import HydraConfig + if HydraConfig.initialized(): + output_dir = HydraConfig.get().runtime.output_dir + if output_dir: + return Path(output_dir).resolve() + except Exception: + pass + return Path.cwd().resolve() + + +_maybe_reexec_without_problematic_ld_preload() + + +def _configure_cuda_runtime(cfg): + """Apply process-level CUDA runtime switches required by this environment.""" + if str(cfg.train.device).startswith('cuda') and bool(cfg.train.get('disable_cudnn', False)): + torch.backends.cudnn.enabled = False + log.warning('⚠️ 已按配置禁用 cuDNN；GPU 卷积将回退到非-cuDNN 实现') + + +def recursive_to_device(data, device): + """ + 递归地将嵌套字典/列表中的张量移动到指定设备。 + + Args: + data: 字典、列表或张量 + device: 目标设备 (例如 'cuda', 'cpu') + + Returns: + 所有张量已移动到指定设备的数据结构 + """ + if isinstance(data, torch.Tensor): + return data.to(device) + elif isinstance(data, dict): + return {k: recursive_to_device(v, device) for k, v in data.items()} + elif isinstance(data, list): + return [recursive_to_device(v, device) for v in data] + return data + + +def resolve_resume_checkpoint(resume_ckpt, checkpoint_dir): + """ + 解析恢复训练用的 checkpoint 路径。 + + Args: + resume_ckpt: 配置中的 resume_ckpt，支持路径或 "auto" + checkpoint_dir: 默认检查点目录 + + Returns: + Path 或 None + """ + if resume_ckpt is None: + return None + + if str(resume_ckpt).lower() != "auto": + return Path(resume_ckpt) + + pattern = re.compile(r"vla_model_step_(\d+)\.pt$") + candidates = [] + for ckpt_path in checkpoint_dir.glob("vla_model_step_*.pt"): + match = pattern.search(ckpt_path.name) + if match: + candidates.append((int(match.group(1)), ckpt_path)) + + if not candidates: + return None + return max(candidates, key=lambda x: x[0])[1] + + +def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_type='cosine', min_lr=0): + """ + 创建带预热的学习率调度器。 + + Args: + optimizer: PyTorch 优化器 + warmup_steps: 预热步数 + max_steps: 总训练步数 + scheduler_type: 预热后的调度器类型 ('cosine' 或 'constant') + min_lr: 最小学习率（用于余弦衰减） + + Returns: + LambdaLR 调度器 + """ + import math + # 在 LambdaLR 修改前捕获初始学习率 + base_lr = optimizer.param_groups[0]['lr'] + min_lr_ratio = min_lr / base_lr if base_lr > 0 else 0.0 + + def lr_lambda(step): + # 预热阶段：从 0 线性增加到 1 + if step < warmup_steps: + return float(step) / float(max(1, warmup_steps)) + + # 预热后阶段 + if scheduler_type == 'cosine': + # 从 1 到 min_lr_ratio 的余弦退火 + progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps)) + cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress)) + return max(min_lr_ratio, cosine_decay) + else: + # 恒定学习率 + return 1.0 + + return LambdaLR(optimizer, lr_lambda) + + +def build_training_optimizer(agent, lr, weight_decay): + """为训练脚本构建优化器，优先复用任意 head 自带的参数分组。""" + trainable_params = [param for param in agent.parameters() if param.requires_grad] + noise_pred_net = getattr(agent, 'noise_pred_net', None) + get_optim_groups = getattr(noise_pred_net, 'get_optim_groups', None) + use_head_groups = callable(get_optim_groups) + + if not use_head_groups: + return AdamW(trainable_params, lr=lr, weight_decay=weight_decay) + + head_groups = [] + grouped_param_ids = set() + for group in get_optim_groups(weight_decay=weight_decay): + params = [param for param in group['params'] if param.requires_grad] + if not params: + continue + normalized_group = dict(group) + normalized_group['params'] = params + head_groups.append(normalized_group) + + for param in params: + param_id = id(param) + if param_id in grouped_param_ids: + raise ValueError('Head optimizer groups contain duplicate parameters') + grouped_param_ids.add(param_id) + + head_trainable_param_ids = { + id(param) for param in noise_pred_net.parameters() if param.requires_grad + } + missing_head_param_ids = head_trainable_param_ids - grouped_param_ids + if missing_head_param_ids: + raise ValueError('Head optimizer groups missed trainable head parameters') + + remaining_params = [ + param for param in trainable_params + if id(param) not in grouped_param_ids + ] + + optim_groups = head_groups + if remaining_params: + optim_groups = optim_groups + [{ + 'params': remaining_params, + 'weight_decay': weight_decay, + }] + grouped_param_ids.update(id(param) for param in remaining_params) + + all_trainable_param_ids = {id(param) for param in trainable_params} + if grouped_param_ids != all_trainable_param_ids: + raise ValueError('Optimizer parameter groups must include each trainable parameter exactly once') + + return AdamW(optim_groups, lr=lr, weight_decay=weight_decay) + + +def _init_swanlab(cfg): + """按需初始化 SwanLab，并在缺少依赖或认证失败时快速失败。""" + if not bool(cfg.train.get('use_swanlab', False)): + return None + + try: + swanlab = importlib.import_module("swanlab") + except ImportError as exc: + raise RuntimeError( + "SwanLab logging is enabled, but the 'swanlab' package could not be imported." + ) from exc + + def _to_plain_config(value): + if isinstance(value, dict): + return {key: _to_plain_config(val) for key, val in value.items()} + if isinstance(value, list): + return [_to_plain_config(item) for item in value] + if isinstance(value, tuple): + return tuple(_to_plain_config(item) for item in value) + + items_method = getattr(value, 'items', None) + if callable(items_method): + try: + return {key: _to_plain_config(val) for key, val in items_method()} + except Exception: + pass + + return value + + swanlab_config = { + key: _to_plain_config(cfg[key]) + for key in ('train', 'data', 'agent') + if key in cfg + } + + init_kwargs = { + 'project': cfg.train.get('swanlab_project', 'roboimi-vla'), + 'config': swanlab_config, + } + run_name = cfg.train.get('swanlab_run_name', None) + if run_name: + init_kwargs['experiment_name'] = run_name + + try: + swanlab.init(**init_kwargs) + except Exception as exc: + raise RuntimeError( + f"SwanLab logging is enabled, but SwanLab init/login failed: {exc}" + ) from exc + + return swanlab + + +def _log_to_swanlab(swanlab_module, payload, step=None): + if swanlab_module is None: + return + try: + swanlab_module.log(payload, step=step) + except Exception as exc: + log.warning(f"SwanLab log failed at step {step}: {exc}") + + +def _log_rollout_trajectory_images_to_swanlab( + swanlab_module, + rollout_stats, + step=None, + context_label: str = 'rollout', +): + if swanlab_module is None or not rollout_stats: + return + + image_factory = getattr(swanlab_module, 'Image', None) + if image_factory is None: + return + + payload = {} + for fallback_episode_index, episode in enumerate(rollout_stats.get('episodes', [])): + if not isinstance(episode, dict): + continue + artifact_paths = episode.get('artifact_paths', {}) + if not isinstance(artifact_paths, dict): + continue + trajectory_image = artifact_paths.get('trajectory_image') + if not trajectory_image: + continue + episode_index = int(episode.get('episode_index', fallback_episode_index)) + caption = f'{context_label} trajectory image - episode {episode_index} (front)' + try: + payload[f'rollout/trajectory_image_episode_{episode_index}'] = image_factory( + str(trajectory_image), + caption=caption, + ) + except Exception as exc: + log.warning( + f"SwanLab rollout trajectory image upload prep failed at step {step}: {exc}" + ) + + if payload: + _log_to_swanlab(swanlab_module, payload, step=step) + + +def _finish_swanlab(swanlab_module): + if swanlab_module is None: + return + try: + swanlab_module.finish() + except Exception as exc: + log.warning(f"SwanLab finish failed: {exc}") + + +def _run_training(cfg: DictConfig): + """ + VLA 训练脚本（ResNet 骨干网络 + Diffusion 策略） + + 该脚本功能： + 1. 从 HDF5 文件加载数据集 + 2. 实例化带 ResNet 视觉编码器的 VLAAgent + 3. 训练基于扩散的动作预测模型 + 4. 定期保存检查点 + """ + + # 打印配置 + print("=" * 80) + print("VLA 训练配置:") + print("=" * 80) + print(OmegaConf.to_yaml(cfg)) + print("=" * 80) + + log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})") + _configure_cuda_runtime(cfg) + swanlab_module = _init_swanlab(cfg) + try: + # 创建检查点目录 + run_output_dir = _resolve_run_output_dir() + checkpoint_dir = run_output_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + default_best_model_path = checkpoint_dir / "vla_model_best.pt" + + # ========================================================================= + # 1. 实例化数据集与 DataLoader + # ========================================================================= + log.info("📦 加载数据集...") + try: + dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224)) + vision_backbone_cfg = cfg.agent.get('vision_backbone', None) + if vision_backbone_cfg is not None and 'dataset_image_resize_shape' in vision_backbone_cfg: + dataset_image_resize_shape = vision_backbone_cfg.get('dataset_image_resize_shape') + dataset = instantiate( + cfg.data, + image_resize_shape=dataset_image_resize_shape, + ) + log.info(f"✅ 数据集加载成功。总样本数: {len(dataset)}") + except Exception as e: + log.error(f"❌ 数据集加载失败: {e}") + raise + + # 训练/验证集划分 + val_split = float(cfg.train.get('val_split', 0.1)) + seed = int(cfg.train.get('seed', 42)) + val_size = int(len(dataset) * val_split) + train_size = len(dataset) - val_size + if val_size > 0: + train_dataset, val_dataset = random_split( + dataset, + [train_size, val_size], + generator=torch.Generator().manual_seed(seed) + ) + log.info(f"✅ 数据集划分: 训练集={train_size}, 验证集={val_size} (验证比例={val_split})") + else: + train_dataset, val_dataset = dataset, None + log.info("✅ 数据集划分: 全部用于训练, 验证集=0 (验证比例=0)") + + train_batch_size = int(cfg.train.batch_size) + train_drop_last = len(train_dataset) >= train_batch_size + if not train_drop_last: + log.warning( + "⚠️ 训练集样本数 (%s) 小于 batch_size (%s)，将保留最后一个不完整批次以避免空训练加载器", + len(train_dataset), + train_batch_size, + ) + + train_loader = DataLoader( + train_dataset, + batch_size=train_batch_size, + shuffle=True, + num_workers=cfg.train.num_workers, + pin_memory=(cfg.train.device != "cpu"), + persistent_workers=False, + drop_last=train_drop_last + ) + + val_loader = None + if val_dataset is not None: + val_loader = DataLoader( + val_dataset, + batch_size=train_batch_size, + shuffle=False, + num_workers=cfg.train.num_workers, + pin_memory=(cfg.train.device != "cpu"), + persistent_workers=False, + drop_last=False + ) + + log.info(f"✅ 训练加载器每轮批次数: {len(train_loader)}") + if val_loader is not None: + log.info(f"✅ 验证加载器每轮批次数: {len(val_loader)}") + + # ========================================================================= + # 2. 加载数据集统计信息（将传递给 agent） + # ========================================================================= + log.info("💾 加载数据集统计信息...") + dataset_stats = None + try: + dataset_dir = cfg.data.get('dataset_dir', 'roboimi/demos/dataset/sim_transfer') + stats_path = Path(dataset_dir) / 'dataset_stats.pkl' + + if stats_path.exists(): + with open(stats_path, 'rb') as f: + stats = pickle.load(f) + + # 扁平化stats字典（嵌套结构→扁平结构）以匹配NormalizationModule的期望格式 + dataset_stats = { + 'action_mean': stats['action_mean'].tolist(), + 'action_std': stats['action_std'].tolist(), + 'action_min': stats['action_min'].tolist(), + 'action_max': stats['action_max'].tolist(), + 'qpos_mean': stats['qpos_mean'].tolist(), + 'qpos_std': stats['qpos_std'].tolist(), + 'qpos_min': stats['qpos_min'].tolist(), + 'qpos_max': stats['qpos_max'].tolist(), + } + log.info(f"✅ 数据集统计信息加载完成 (归一化: {cfg.agent.normalization_type})") + else: + log.warning(f"⚠️ 统计文件未找到: {stats_path}") + log.warning("⚠️ 推理时动作将无法反归一化！") + + except Exception as e: + log.warning(f"⚠️ 统计信息加载失败: {e}") + log.warning("⚠️ 训练将继续，但推理可能无法正常工作") + + # ========================================================================= + # 3. 实例化 VLA Agent + # ========================================================================= + log.info("🤖 初始化 VLA Agent...") + try: + # 将 dataset_stats 和 normalization_type 传递给 agent + agent = instantiate(cfg.agent, dataset_stats=dataset_stats) + agent.to(cfg.train.device) + agent.train() + log.info(f"✅ Agent 初始化完成并已移至 {cfg.train.device}") + + # 统计参数量 + total_params = sum(p.numel() for p in agent.parameters()) + trainable_params = sum(p.numel() for p in agent.parameters() if p.requires_grad) + log.info(f"📊 总参数量: {total_params:,}") + log.info(f"📊 可训练参数量: {trainable_params:,}") + + except Exception as e: + log.error(f"❌ Agent 初始化失败: {e}") + raise + + # ========================================================================= + # 3.1 从预训练 checkpoint 加载权重（微调） + # ========================================================================= + pretrained_ckpt = cfg.train.get('pretrained_ckpt', None) + if pretrained_ckpt is not None: + ckpt_path = Path(pretrained_ckpt) + if ckpt_path.exists(): + log.info(f"🔄 [Finetune] 从预训练 checkpoint 加载权重: {ckpt_path}") + try: + checkpoint = torch.load(ckpt_path, map_location=cfg.train.device) + + # 只加载模型权重（不加载 optimizer、scheduler） + missing_keys, unexpected_keys = agent.load_state_dict( + checkpoint['model_state_dict'], + strict=False # 允许部分加载（结构不完全匹配时） + ) + + log.info(f"✅ [Finetune] 模型权重加载成功") + + if missing_keys: + log.warning(f"⚠️ [Finetune] 缺少的键 ({len(missing_keys)} 个): {missing_keys[:5]}...") + if unexpected_keys: + log.warning(f"⚠️ [Finetune] 多余的键 ({len(unexpected_keys)} 个): {unexpected_keys[:5]}...") + + log.info(f"📊 [Finetune] 预训练信息: 步骤={checkpoint.get('step', 'N/A')}, 损失={checkpoint.get('loss', 'N/A')}") + log.info(f"📈 [Finetune] 使用新的训练配置（lr={cfg.train.lr}, max_steps={cfg.train.max_steps}）") + + except Exception as e: + log.error(f"❌ [Finetune] 加载 checkpoint 失败: {e}") + log.warning("⚠️ 将从头开始训练") + else: + log.error(f"❌ [Finetune] Checkpoint 文件不存在: {ckpt_path}") + log.warning("⚠️ 将从头开始训练") + + # ========================================================================= + # 4. 设置优化器与学习率调度器 + # ========================================================================= + weight_decay = float(cfg.train.get('weight_decay', 1e-5)) + grad_clip = float(cfg.train.get('grad_clip', 1.0)) + + optimizer = build_training_optimizer(agent, lr=cfg.train.lr, weight_decay=weight_decay) + log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr}, weight_decay={weight_decay})") + + # 设置带预热的学習率调度器 + warmup_steps = int(cfg.train.get('warmup_steps', 500)) + scheduler_type = cfg.train.get('scheduler_type', 'cosine') + min_lr = float(cfg.train.get('min_lr', 1e-6)) + + scheduler = get_lr_schedule_with_warmup( + optimizer, + warmup_steps=warmup_steps, + max_steps=cfg.train.max_steps, + scheduler_type=scheduler_type, + min_lr=min_lr + ) + log.info(f"📈 学习率调度器: {scheduler_type}，{warmup_steps} 步预热 (最小学习率={min_lr})") + + # ========================================================================= + # 4.1 断点续训（恢复模型、优化器、调度器、步数） + # ========================================================================= + def extract_checkpoint_metric_baseline(checkpoint): + checkpoint_loss = checkpoint.get('loss', None) + checkpoint_val_loss = checkpoint.get('val_loss', None) + checkpoint_rollout_reward = checkpoint.get('rollout_avg_reward', None) + + baseline_loss = float('inf') + baseline_rollout_reward = float('-inf') + if checkpoint_rollout_reward is not None: + baseline_rollout_reward = float(checkpoint_rollout_reward) + if checkpoint_val_loss is not None: + baseline_loss = float(checkpoint_val_loss) + elif checkpoint_loss is not None: + baseline_loss = float(checkpoint_loss) + return baseline_loss, baseline_rollout_reward + + start_step = 0 + resume_loss = None + resume_best_loss = float('inf') + resume_best_rollout_reward = float('-inf') + best_model_path = None + + resume_ckpt = cfg.train.get('resume_ckpt', None) + resume_path = resolve_resume_checkpoint(resume_ckpt, checkpoint_dir) + if resume_ckpt is not None: + if pretrained_ckpt is not None: + log.warning("⚠️ [Resume] 同时设置了 pretrained_ckpt 与 resume_ckpt，将优先使用 resume_ckpt 进行断点续训") + if resume_path is None: + log.warning("⚠️ [Resume] 未找到可恢复的 checkpoint，将从头开始训练") + elif not resume_path.exists(): + log.error(f"❌ [Resume] Checkpoint 文件不存在: {resume_path}") + log.warning("⚠️ 将从头开始训练") + else: + log.info(f"🔄 [Resume] 从 checkpoint 恢复训练: {resume_path}") + try: + checkpoint = torch.load(resume_path, map_location=cfg.train.device) + + agent.load_state_dict(checkpoint['model_state_dict'], strict=True) + optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + scheduler.load_state_dict(checkpoint['scheduler_state_dict']) + + resume_step = int(checkpoint['step']) + start_step = resume_step + 1 + + loaded_loss = checkpoint.get('loss', None) + resume_loss = float(loaded_loss) if loaded_loss is not None else None + resume_best_loss, resume_best_rollout_reward = extract_checkpoint_metric_baseline(checkpoint) + if ( + resume_best_rollout_reward != float('-inf') + or resume_best_loss != float('inf') + ): + best_model_path = resume_path + + if default_best_model_path.exists(): + try: + best_checkpoint = torch.load(default_best_model_path, map_location=cfg.train.device) + _, best_checkpoint_rollout_reward = ( + extract_checkpoint_metric_baseline(best_checkpoint) + ) + if best_checkpoint_rollout_reward != float('-inf'): + resume_best_rollout_reward = best_checkpoint_rollout_reward + best_model_path = default_best_model_path + log.info( + "📈 [Resume] 从最佳 checkpoint 恢复最佳 rollout 基线: %s", + default_best_model_path, + ) + except Exception as e: + log.warning( + f"⚠️ [Resume] 读取最佳 checkpoint 失败，将回退到恢复 checkpoint 的验证基线: {e}" + ) + + log.info(f"✅ [Resume] 恢复成功: 上次步骤={resume_step}, 本次从步骤 {start_step} 开始") + log.info(f"📈 [Resume] 当前学习率: {optimizer.param_groups[0]['lr']:.2e}") + except Exception as e: + log.error(f"❌ [Resume] 恢复失败: {e}") + log.warning("⚠️ 将从头开始训练") + start_step = 0 + resume_loss = None + resume_best_loss = float('inf') + resume_best_rollout_reward = float('-inf') + + # ========================================================================= + # 5. 训练循环 + # ========================================================================= + log.info("🏋️ 开始训练循环...") + + def build_agent_input(batch_data): + """构建 agent 输入格式""" + images = {} + # SimpleRobotDataset 返回 observation.{cam_name} 格式 + for cam_name in cfg.data.camera_names: + key = f"observation.{cam_name}" + if key in batch_data: + images[cam_name] = batch_data[key] + + return { + 'images': images, + 'qpos': batch_data['observation.state'], # SimpleRobotDataset 使用 observation.state + 'action': batch_data['action'], + 'action_is_pad': batch_data.get('action_is_pad', None) # 传递padding mask + } + + def save_checkpoint(checkpoint_path: Path, step: int, loss_value, val_loss=None, rollout_avg_reward=None): + agent_stats = agent.get_normalization_stats() + torch.save({ + 'step': step, + 'model_state_dict': agent.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'scheduler_state_dict': scheduler.state_dict(), + 'loss': loss_value, + 'val_loss': val_loss, + 'rollout_avg_reward': rollout_avg_reward, + 'dataset_stats': agent_stats, # 保存agent的统计信息 + 'current_lr': optimizer.param_groups[0]['lr'], + }, checkpoint_path) + return checkpoint_path + + def run_validation(): + """运行验证""" + if val_loader is None: + return None + agent.eval() + + # 设置确定性种子以获得可重现的损失 + # 这确保验证损失在不同步骤之间可比较 + torch.manual_seed(42) + if torch.cuda.is_available(): + torch.cuda.manual_seed(42) + + total_loss = 0.0 + num_batches = 0 + with torch.no_grad(): + for val_batch in val_loader: + val_batch = recursive_to_device(val_batch, cfg.train.device) + val_input = build_agent_input(val_batch) + val_loss = agent.compute_loss(val_input) + total_loss += val_loss.item() + num_batches += 1 + agent.train() + return total_loss / max(num_batches, 1) + + def run_rollout_validation(checkpoint_path: Path): + from roboimi.demos.vla_scripts import eval_vla + + rollout_cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=False)) + rollout_cfg.eval.ckpt_path = str(checkpoint_path) + rollout_cfg.eval.num_episodes = int(cfg.train.get('rollout_num_episodes', 1)) + rollout_cfg.eval.headless = True + rollout_cfg.eval.device = 'cpu' + rollout_cfg.eval.verbose_action = False + rollout_cfg.eval.record_video = False + rollout_cfg.eval.save_trajectory_image = True + rollout_cfg.eval.trajectory_image_camera_name = 'front' + rollout_cfg.eval.save_summary_json = True + rollout_cfg.eval.artifact_dir = str( + (run_output_dir / 'rollout_artifacts' / checkpoint_path.stem).resolve() + ) + + log.info( + "🎯 开始 checkpoint rollout 验证: %s (episodes=%s, headless=True)", + checkpoint_path, + rollout_cfg.eval.num_episodes, + ) + return eval_vla._run_eval(rollout_cfg) + + def run_checkpoint_rollout_validation(checkpoint_path: Path): + if not bool(cfg.train.get('rollout_validate_on_checkpoint', False)): + return None + return run_rollout_validation(checkpoint_path) + + data_iter = iter(train_loader) + pbar = tqdm(range(start_step, cfg.train.max_steps), desc="训练中", ncols=100) + + steps_per_epoch = len(train_loader) + rollout_val_freq_epochs = int(cfg.train.get('rollout_val_freq_epochs', 0) or 0) + rollout_validation_enabled = rollout_val_freq_epochs > 0 + best_loss = resume_best_loss + best_rollout_reward = resume_best_rollout_reward + last_loss = resume_loss + + if start_step >= cfg.train.max_steps: + log.warning( + f"⚠️ [Resume] start_step={start_step} 已达到/超过 max_steps={cfg.train.max_steps}，跳过训练循环" + ) + + for step in pbar: + try: + batch = next(data_iter) + except StopIteration: + # 轮次结束时重启迭代器 + data_iter = iter(train_loader) + batch = next(data_iter) + + # ===================================================================== + # 将批次移至设备 + # ===================================================================== + batch = recursive_to_device(batch, cfg.train.device) + + # ===================================================================== + # 准备 agent 输入 + # ===================================================================== + # 数据集返回: {action, qpos, image_, ...} + # Agent 期望: {images: dict, qpos: tensor, action: tensor} + + # 准备 agent 输入 + agent_input = build_agent_input(batch) + + # ===================================================================== + # 前向传播与损失计算 + # ===================================================================== + try: + loss = agent.compute_loss(agent_input) + except Exception as e: + log.error(f"❌ 步骤 {step} 前向传播失败: {e}") + raise + + last_loss = loss.item() + + # ===================================================================== + # 反向传播与优化 + # ===================================================================== + optimizer.zero_grad() + loss.backward() + + # 梯度裁剪以稳定训练 + torch.nn.utils.clip_grad_norm_(agent.parameters(), max_norm=grad_clip) + + optimizer.step() + scheduler.step() + + # ===================================================================== + # 日志记录 + # ===================================================================== + if step % cfg.train.log_freq == 0: + current_lr = optimizer.param_groups[0]['lr'] + best_loss_to_log = best_loss if best_loss != float('inf') else loss.item() + pbar.set_postfix({ + "loss": f"{loss.item():.4f}", + "lr": f"{current_lr:.2e}", + "best_loss": f"{best_loss_to_log:.4f}" + }) + log.info(f"步骤 {step}/{cfg.train.max_steps} | 损失: {loss.item():.4f} | 学习率: {current_lr:.2e}") + _log_to_swanlab( + swanlab_module, + { + 'train/loss': loss.item(), + 'train/lr': current_lr, + 'train/best_loss': best_loss_to_log, + 'train/step': step, + }, + step=step, + ) + + # ===================================================================== + # 检查点保存与验证 + # ===================================================================== + checkpoint_path = None + val_loss = None + if step > 0 and step % cfg.train.save_freq == 0: + # 运行验证 + val_loss = run_validation() + if val_loss is not None: + log.info(f"步骤 {step}/{cfg.train.max_steps} | 验证损失: {val_loss:.4f}") + _log_to_swanlab( + swanlab_module, + {'val/loss': val_loss}, + step=step, + ) + + checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt" + save_checkpoint( + checkpoint_path, + step, + loss.item(), + val_loss=val_loss, + ) + log.info(f"💾 检查点已保存: {checkpoint_path}") + + # 在首次拿到 rollout 平均奖励之前，使用损失作为最佳模型回退指标 + if best_rollout_reward == float('-inf'): + eval_loss = val_loss if val_loss is not None else loss.item() + if eval_loss < best_loss: + best_loss = eval_loss + best_model_path = default_best_model_path + save_checkpoint( + best_model_path, + step, + loss.item(), + val_loss=val_loss, + ) + log.info(f"🌟 最佳模型已更新: {best_model_path} (验证损失: {best_loss:.4f})") + + checkpoint_rollout_stats = run_checkpoint_rollout_validation(checkpoint_path) + checkpoint_rollout_avg_reward = ( + checkpoint_rollout_stats.get('avg_reward') + if checkpoint_rollout_stats is not None else None + ) + if checkpoint_rollout_avg_reward is not None: + log.info( + f"步骤 {step}/{cfg.train.max_steps} | checkpoint rollout 平均奖励: " + f"{checkpoint_rollout_avg_reward:.4f}" + ) + _log_to_swanlab( + swanlab_module, + {'rollout/avg_reward': checkpoint_rollout_avg_reward}, + step=step, + ) + if checkpoint_rollout_avg_reward > best_rollout_reward: + best_rollout_reward = checkpoint_rollout_avg_reward + best_model_path = default_best_model_path + save_checkpoint( + best_model_path, + step, + loss.item(), + val_loss=val_loss, + rollout_avg_reward=checkpoint_rollout_avg_reward, + ) + log.info( + f"🌟 最佳模型已更新: {best_model_path} " + f"(checkpoint rollout 平均奖励: {best_rollout_reward:.4f})" + ) + + completed_steps = step + 1 + completed_epoch = ( + completed_steps // steps_per_epoch + if steps_per_epoch > 0 else 0 + ) + should_run_epoch_rollout = ( + rollout_validation_enabled + and steps_per_epoch > 0 + and completed_steps % steps_per_epoch == 0 + and completed_epoch > 0 + and completed_epoch % rollout_val_freq_epochs == 0 + ) + if should_run_epoch_rollout: + if checkpoint_path is None: + checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt" + save_checkpoint( + checkpoint_path, + step, + loss.item(), + val_loss=val_loss, + ) + log.info(f"💾 Epoch rollout 验证前检查点已保存: {checkpoint_path}") + + rollout_stats = run_rollout_validation(checkpoint_path) + rollout_avg_reward = ( + rollout_stats.get('avg_reward') + if rollout_stats is not None else None + ) + if rollout_avg_reward is not None: + log.info( + f"步骤 {step}/{cfg.train.max_steps} | Epoch {completed_epoch} " + f"rollout 平均奖励: {rollout_avg_reward:.4f}" + ) + _log_to_swanlab( + swanlab_module, + { + 'rollout/avg_reward': rollout_avg_reward, + 'rollout/epoch': completed_epoch, + }, + step=step, + ) + _log_rollout_trajectory_images_to_swanlab( + swanlab_module, + rollout_stats, + step=step, + context_label=f'epoch {completed_epoch} rollout', + ) + if rollout_avg_reward > best_rollout_reward: + best_rollout_reward = rollout_avg_reward + best_model_path = default_best_model_path + save_checkpoint( + best_model_path, + step, + loss.item(), + val_loss=val_loss, + rollout_avg_reward=rollout_avg_reward, + ) + log.info( + f"🌟 最佳模型已更新: {best_model_path} " + f"(Epoch {completed_epoch} rollout 平均奖励: {best_rollout_reward:.4f})" + ) + + # ========================================================================= + # 6. 保存最终模型 + # ========================================================================= + final_model_path = checkpoint_dir / "vla_model_final.pt" + save_checkpoint( + final_model_path, + cfg.train.max_steps, + last_loss, + ) + log.info(f"💾 最终模型已保存: {final_model_path}") + _log_to_swanlab( + swanlab_module, + { + 'final/checkpoint_path': str(final_model_path), + 'final/best_checkpoint_path': ( + str(best_model_path) if best_model_path is not None else '' + ), + }, + step=cfg.train.max_steps, + ) + + log.info("✅ 训练成功完成!") + if last_loss is not None: + log.info(f"📊 最终损失: {last_loss:.4f}") + else: + log.info("📊 最终损失: N/A（未执行训练步）") + if best_rollout_reward != float('-inf'): + log.info(f"📊 最佳 rollout 平均奖励: {best_rollout_reward:.4f}") + elif best_loss != float('inf'): + log.info(f"📊 最佳损失: {best_loss:.4f}") + else: + log.info("📊 最佳验证指标: N/A（无有效 rollout/验证损失）") + finally: + _finish_swanlab(swanlab_module) + + +@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config") +def main(cfg: DictConfig): + _run_training(cfg) + + +if __name__ == "__main__": + main() diff --git a/roboimi/detr/LICENSE b/roboimi/detr/LICENSE deleted file mode 100644 index b1395e9..0000000 --- a/roboimi/detr/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2020 - present, Facebook, Inc - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/roboimi/detr/README.md b/roboimi/detr/README.md deleted file mode 100644 index 500b1b8..0000000 --- a/roboimi/detr/README.md +++ /dev/null @@ -1,9 +0,0 @@ -This part of the codebase is modified from DETR https://github.com/facebookresearch/detr under APACHE 2.0. - - @article{Carion2020EndtoEndOD, - title={End-to-End Object Detection with Transformers}, - author={Nicolas Carion and Francisco Massa and Gabriel Synnaeve and Nicolas Usunier and Alexander Kirillov and Sergey Zagoruyko}, - journal={ArXiv}, - year={2020}, - volume={abs/2005.12872} - } \ No newline at end of file diff --git a/roboimi/detr/main.py b/roboimi/detr/main.py deleted file mode 100644 index 4891049..0000000 --- a/roboimi/detr/main.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -import argparse -from pathlib import Path - -import numpy as np -import torch -from .models import build_ACT_model, build_CNNMLP_model - - -def get_args_parser(): - parser = argparse.ArgumentParser('Set transformer detector', add_help=False) - parser.add_argument('--lr', default=1e-4, type=float) # will be overridden - parser.add_argument('--lr_backbone', default=1e-5, type=float) # will be overridden - parser.add_argument('--batch_size', default=2, type=int) # not used - parser.add_argument('--weight_decay', default=1e-4, type=float) - parser.add_argument('--epochs', default=300, type=int) # not used - parser.add_argument('--lr_drop', default=200, type=int) # not used - parser.add_argument('--clip_max_norm', default=0.1, type=float, # not used - help='gradient clipping max norm') - parser.add_argument('--qpos_noise_std', action='store', default=0, type=float, help='lr', required=False) - - # Model parameters - # * Backbone - parser.add_argument('--backbone', default='resnet18', type=str, # will be overridden - help="Name of the convolutional backbone to use") - parser.add_argument('--dilation', action='store_true', - help="If true, we replace stride with dilation in the last convolutional block (DC5)") - parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), - help="Type of positional embedding to use on top of the image features") - parser.add_argument('--camera_names', default=[], type=list, # will be overridden - help="A list of camera names") - - # * Transformer - parser.add_argument('--enc_layers', default=4, type=int, # will be overridden - help="Number of encoding layers in the transformer") - parser.add_argument('--dec_layers', default=6, type=int, # will be overridden - help="Number of decoding layers in the transformer") - parser.add_argument('--dim_feedforward', default=2048, type=int, # will be overridden - help="Intermediate size of the feedforward layers in the transformer blocks") - parser.add_argument('--hidden_dim', default=256, type=int, # will be overridden - help="Size of the embeddings (dimension of the transformer)") - parser.add_argument('--dropout', default=0.1, type=float, - help="Dropout applied in the transformer") - parser.add_argument('--nheads', default=8, type=int, # will be overridden - help="Number of attention heads inside the transformer's attentions") - parser.add_argument('--num_queries', default=400, type=int, # will be overridden - help="Number of query slots") - parser.add_argument('--pre_norm', action='store_true') - parser.add_argument('--state_dim', default=14, type=int) - parser.add_argument('--action_dim', default=14, type=int) - - - # * Segmentation - parser.add_argument('--masks', action='store_true', - help="Train segmentation head if the flag is provided") - - - - return parser - - -def build_ACT_model_and_optimizer(args_override): - parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()]) - args = parser.parse_args() - - for k, v in args_override.items(): - setattr(args, k, v) - - model = build_ACT_model(args) - model.cuda() - - param_dicts = [ - {"params": [p for n, p in model.named_parameters() if "backbone" not in n and p.requires_grad]}, - { - "params": [p for n, p in model.named_parameters() if "backbone" in n and p.requires_grad], - "lr": args.lr_backbone, - }, - ] - optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, - weight_decay=args.weight_decay) - - return model, optimizer - - -def build_CNNMLP_model_and_optimizer(args_override): - parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()]) - args = parser.parse_args() - - for k, v in args_override.items(): - setattr(args, k, v) - - model = build_CNNMLP_model(args) - model.cuda() - - param_dicts = [ - {"params": [p for n, p in model.named_parameters() if "backbone" not in n and p.requires_grad]}, - { - "params": [p for n, p in model.named_parameters() if "backbone" in n and p.requires_grad], - "lr": args.lr_backbone, - }, - ] - optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, - weight_decay=args.weight_decay) - - return model, optimizer - diff --git a/roboimi/detr/models/__init__.py b/roboimi/detr/models/__init__.py deleted file mode 100644 index cc78db1..0000000 --- a/roboimi/detr/models/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -from .detr_vae import build as build_vae -from .detr_vae import build_cnnmlp as build_cnnmlp - -def build_ACT_model(args): - return build_vae(args) - -def build_CNNMLP_model(args): - return build_cnnmlp(args) \ No newline at end of file diff --git a/roboimi/detr/models/detr_vae.py b/roboimi/detr/models/detr_vae.py deleted file mode 100644 index afcdc5d..0000000 --- a/roboimi/detr/models/detr_vae.py +++ /dev/null @@ -1,300 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -DETR model and criterion classes. -""" -import torch -from torch import nn -from torch.autograd import Variable -from .backbone import build_backbone -from .transformer import build_transformer, TransformerEncoder, TransformerEncoderLayer - -import numpy as np - - -def reparametrize(mu, logvar): - std = logvar.div(2).exp() - eps = Variable(std.data.new(std.size()).normal_()) - return mu + std * eps - - -def get_sinusoid_encoding_table(n_position, d_hid): - def get_position_angle_vec(position): - return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] - - sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) - sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i - sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 - - return torch.FloatTensor(sinusoid_table).unsqueeze(0) - - -class DETRVAE(nn.Module): - """ This is the DETR module that performs object detection """ - def __init__(self, backbones, transformer, encoder, state_dim, action_dim, num_queries, camera_names): - """ Initializes the model. - Parameters: - backbones: torch module of the backbone to be used. See backbone.py - transformer: torch module of the transformer architecture. See transformer.py - state_dim: robot state dimension of the environment - num_queries: number of object queries, ie detection slot. This is the maximal number of objects - DETR can detect in a single image. For COCO, we recommend 100 queries. - aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. - """ - super().__init__() - self.num_queries = num_queries - self.camera_names = camera_names - self.transformer = transformer - self.encoder = encoder - hidden_dim = transformer.d_model - self.action_head = nn.Linear(hidden_dim, action_dim) - self.is_pad_head = nn.Linear(hidden_dim, 1) - self.query_embed = nn.Embedding(num_queries, hidden_dim) - if backbones is not None: - self.input_proj = nn.Conv2d(backbones[0].num_channels, hidden_dim, kernel_size=1) - self.backbones = nn.ModuleList(backbones) - self.input_proj_robot_state = nn.Linear(state_dim, hidden_dim) - else: - raise NotImplementedError - # input_dim = 14 + 7 # robot_state + env_state - # self.input_proj_robot_state = nn.Linear(state_dim, hidden_dim) - # self.input_proj_env_state = nn.Linear(7, hidden_dim) - # self.pos = torch.nn.Embedding(2, hidden_dim) - # self.backbones = None - - # encoder extra parameters - self.latent_dim = 32 # final size of latent z # TODO tune - self.cls_embed = nn.Embedding(1, hidden_dim) # extra cls token embedding - self.encoder_action_proj = nn.Linear(action_dim, hidden_dim) # project action to embedding - self.encoder_joint_proj = nn.Linear(state_dim, hidden_dim) # project qpos to embedding - self.latent_proj = nn.Linear(hidden_dim, self.latent_dim*2) # project hidden state to latent std, var - self.register_buffer('pos_table', get_sinusoid_encoding_table(1+1+num_queries, hidden_dim)) # [CLS], qpos, a_seq - - # decoder extra parameters - self.latent_out_proj = nn.Linear(self.latent_dim, hidden_dim) # project latent sample to embedding - self.additional_pos_embed = nn.Embedding(2, hidden_dim) # learned position embedding for proprio and latent - - def forward(self, qpos, image, env_state, actions=None, is_pad=None): - """ - qpos: batch, qpos_dim - image: batch, num_cam, channel, height, width - env_state: None - actions: batch, seq, action_dim - """ - is_training = actions is not None # train or val - bs, _ = qpos.shape - ### Obtain latent z from action sequence - if is_training: - # project action sequence to embedding dim, and concat with a CLS token - action_embed = self.encoder_action_proj(actions) # (bs, seq, hidden_dim) - qpos_embed = self.encoder_joint_proj(qpos) # (bs, hidden_dim) - qpos_embed = torch.unsqueeze(qpos_embed, axis=1) # (bs, 1, hidden_dim) - cls_embed = self.cls_embed.weight # (1, hidden_dim) - cls_embed = torch.unsqueeze(cls_embed, axis=0).repeat(bs, 1, 1) # (bs, 1, hidden_dim) - encoder_input = torch.cat([cls_embed, qpos_embed, action_embed], axis=1) # (bs, seq+1, hidden_dim) - encoder_input = encoder_input.permute(1, 0, 2) # (seq+1, bs, hidden_dim) - # do not mask cls token - cls_joint_is_pad = torch.full((bs, 2), False).to(qpos.device) # False: not a padding - is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1) # (bs, seq+1) - # obtain position embedding - pos_embed = self.pos_table.clone().detach() - pos_embed = pos_embed.permute(1, 0, 2) # (seq+1, 1, hidden_dim) - # query model - encoder_output = self.encoder(encoder_input, pos=pos_embed, src_key_padding_mask=is_pad) - encoder_output = encoder_output[0] # take cls output only - latent_info = self.latent_proj(encoder_output) - mu = latent_info[:, :self.latent_dim] - logvar = latent_info[:, self.latent_dim:] - latent_sample = reparametrize(mu, logvar) - latent_input = self.latent_out_proj(latent_sample) - else: - mu = logvar = None - latent_sample = torch.zeros([bs, self.latent_dim], dtype=torch.float32).to(qpos.device) - latent_input = self.latent_out_proj(latent_sample) - - if self.backbones is not None: - # Image observation features and position embeddings - all_cam_features = [] - all_cam_pos = [] - - - - - # print(f"Image shape: {image.shape}, Number of cameras: {len(self.camera_names)}") - - - for cam_id, cam_name in enumerate(self.camera_names): - # features, pos = self.backbones[0](image[:, cam_id]) # HARDCODED - features, pos = self.backbones[cam_id](image[:, cam_id]) - features = features[0] # take the last layer feature - pos = pos[0] - all_cam_features.append(self.input_proj(features)) - all_cam_pos.append(pos) - - - - - - - - - - - - # proprioception features - proprio_input = self.input_proj_robot_state(qpos) - # fold camera dimension into width dimension - src = torch.cat(all_cam_features, axis=3) - pos = torch.cat(all_cam_pos, axis=3) - hs = self.transformer(src, None, self.query_embed.weight, pos, latent_input, proprio_input, self.additional_pos_embed.weight)[0] - else: - qpos = self.input_proj_robot_state(qpos) - env_state = self.input_proj_env_state(env_state) - transformer_input = torch.cat([qpos, env_state], axis=1) # seq length = 2 - hs = self.transformer(transformer_input, None, self.query_embed.weight, self.pos.weight)[0] - a_hat = self.action_head(hs) - is_pad_hat = self.is_pad_head(hs) - return a_hat, is_pad_hat, [mu, logvar] - - - -class CNNMLP(nn.Module): - def __init__(self, backbones, state_dim, camera_names): - """ Initializes the model. - Parameters: - backbones: torch module of the backbone to be used. See backbone.py - transformer: torch module of the transformer architecture. See transformer.py - state_dim: robot state dimension of the environment - num_queries: number of object queries, ie detection slot. This is the maximal number of objects - DETR can detect in a single image. For COCO, we recommend 100 queries. - aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. - """ - super().__init__() - self.camera_names = camera_names - self.action_head = nn.Linear(1000, state_dim) # TODO add more - if backbones is not None: - self.backbones = nn.ModuleList(backbones) - backbone_down_projs = [] - for backbone in backbones: - down_proj = nn.Sequential( - nn.Conv2d(backbone.num_channels, 128, kernel_size=5), - nn.Conv2d(128, 64, kernel_size=5), - nn.Conv2d(64, 32, kernel_size=5) - ) - backbone_down_projs.append(down_proj) - self.backbone_down_projs = nn.ModuleList(backbone_down_projs) - - mlp_in_dim = 768 * len(backbones) + 14 - self.mlp = mlp(input_dim=mlp_in_dim, hidden_dim=1024, output_dim=14, hidden_depth=2) - else: - raise NotImplementedError - - def forward(self, qpos, image, env_state, actions=None): - """ - qpos: batch, qpos_dim - image: batch, num_cam, channel, height, width - env_state: None - actions: batch, seq, action_dim - """ - is_training = actions is not None # train or val - bs, _ = qpos.shape - # Image observation features and position embeddings - all_cam_features = [] - for cam_id, cam_name in enumerate(self.camera_names): - features, pos = self.backbones[cam_id](image[:, cam_id]) - features = features[0] # take the last layer feature - pos = pos[0] # not used - all_cam_features.append(self.backbone_down_projs[cam_id](features)) - # flatten everything - flattened_features = [] - for cam_feature in all_cam_features: - flattened_features.append(cam_feature.reshape([bs, -1])) - flattened_features = torch.cat(flattened_features, axis=1) # 768 each - features = torch.cat([flattened_features, qpos], axis=1) # qpos: 14 - a_hat = self.mlp(features) - return a_hat - - -def mlp(input_dim, hidden_dim, output_dim, hidden_depth): - if hidden_depth == 0: - mods = [nn.Linear(input_dim, output_dim)] - else: - mods = [nn.Linear(input_dim, hidden_dim), nn.ReLU(inplace=True)] - for i in range(hidden_depth - 1): - mods += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU(inplace=True)] - mods.append(nn.Linear(hidden_dim, output_dim)) - trunk = nn.Sequential(*mods) - return trunk - - -def build_encoder(args): - d_model = args.hidden_dim # 256 - dropout = args.dropout # 0.1 - nhead = args.nheads # 8 - dim_feedforward = args.dim_feedforward # 2048 - num_encoder_layers = args.enc_layers # 4 # TODO shared with VAE decoder - normalize_before = args.pre_norm # False - activation = "relu" - - encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, - dropout, activation, normalize_before) - encoder_norm = nn.LayerNorm(d_model) if normalize_before else None - encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) - - return encoder - - -def build(args): - state_dim = args.state_dim - action_dim = args.action_dim - - # From state - # backbone = None # from state for now, no need for conv nets - # From image - backbones = [] - # backbone = build_backbone(args) - # backbones.append(backbone) - for _ in args.camera_names: - backbone = build_backbone(args) - backbones.append(backbone) - - transformer = build_transformer(args) - - encoder = build_encoder(args) - - model = DETRVAE( - backbones, - transformer, - encoder, - state_dim=state_dim, - action_dim=action_dim, - num_queries=args.num_queries, - camera_names=args.camera_names, - ) - - n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) - print("number of parameters: %.2fM" % (n_parameters/1e6,)) - - return model - -def build_cnnmlp(args): - state_dim = 14 # TODO hardcode - - # From state - # backbone = None # from state for now, no need for conv nets - # From image - backbones = [] - for _ in args.camera_names: - backbone = build_backbone(args) - backbones.append(backbone) - - model = CNNMLP( - backbones, - state_dim=state_dim, - camera_names=args.camera_names, - ) - - n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) - print("number of parameters: %.2fM" % (n_parameters/1e6,)) - - return model - diff --git a/roboimi/detr/models/transformer.py b/roboimi/detr/models/transformer.py deleted file mode 100644 index 2306ab2..0000000 --- a/roboimi/detr/models/transformer.py +++ /dev/null @@ -1,312 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -DETR Transformer class. - -Copy-paste from torch.nn.Transformer with modifications: - * positional encodings are passed in MHattention - * extra LN at the end of encoder is removed - * decoder returns a stack of activations from all decoding layers -""" -import copy -from typing import Optional, List - -import torch -import torch.nn.functional as F -from torch import nn, Tensor - - -class Transformer(nn.Module): - - def __init__(self, d_model=512, nhead=8, num_encoder_layers=6, - num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, - activation="relu", normalize_before=False, - return_intermediate_dec=False): - super().__init__() - - encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, - dropout, activation, normalize_before) - encoder_norm = nn.LayerNorm(d_model) if normalize_before else None - self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) - - decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, - dropout, activation, normalize_before) - decoder_norm = nn.LayerNorm(d_model) - self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, - return_intermediate=return_intermediate_dec) - - self._reset_parameters() - - self.d_model = d_model - self.nhead = nhead - - def _reset_parameters(self): - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward(self, src, mask, query_embed, pos_embed, latent_input=None, proprio_input=None, additional_pos_embed=None): - # TODO flatten only when input has H and W - if len(src.shape) == 4: # has H and W - # flatten NxCxHxW to HWxNxC - bs, c, h, w = src.shape - src = src.flatten(2).permute(2, 0, 1) - pos_embed = pos_embed.flatten(2).permute(2, 0, 1).repeat(1, bs, 1) - query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) - # mask = mask.flatten(1) - - additional_pos_embed = additional_pos_embed.unsqueeze(1).repeat(1, bs, 1) # seq, bs, dim - pos_embed = torch.cat([additional_pos_embed, pos_embed], axis=0) - - addition_input = torch.stack([latent_input, proprio_input], axis=0) - src = torch.cat([addition_input, src], axis=0) - else: - assert len(src.shape) == 3 - # flatten NxHWxC to HWxNxC - bs, hw, c = src.shape - src = src.permute(1, 0, 2) - pos_embed = pos_embed.unsqueeze(1).repeat(1, bs, 1) - query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) - - tgt = torch.zeros_like(query_embed) - memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) - hs = self.decoder(tgt, memory, memory_key_padding_mask=mask, - pos=pos_embed, query_pos=query_embed) - hs = hs.transpose(1, 2) - return hs - -class TransformerEncoder(nn.Module): - - def __init__(self, encoder_layer, num_layers, norm=None): - super().__init__() - self.layers = _get_clones(encoder_layer, num_layers) - self.num_layers = num_layers - self.norm = norm - - def forward(self, src, - mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - output = src - - for layer in self.layers: - output = layer(output, src_mask=mask, - src_key_padding_mask=src_key_padding_mask, pos=pos) - - if self.norm is not None: - output = self.norm(output) - - return output - - -class TransformerDecoder(nn.Module): - - def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): - super().__init__() - self.layers = _get_clones(decoder_layer, num_layers) - self.num_layers = num_layers - self.norm = norm - self.return_intermediate = return_intermediate - - def forward(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - output = tgt - - intermediate = [] - - for layer in self.layers: - output = layer(output, memory, tgt_mask=tgt_mask, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask, - pos=pos, query_pos=query_pos) - if self.return_intermediate: - intermediate.append(self.norm(output)) - - if self.norm is not None: - output = self.norm(output) - if self.return_intermediate: - intermediate.pop() - intermediate.append(output) - - if self.return_intermediate: - return torch.stack(intermediate) - - return output.unsqueeze(0) - - -class TransformerEncoderLayer(nn.Module): - - def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, - activation="relu", normalize_before=False): - super().__init__() - self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) - # Implementation of Feedforward model - self.linear1 = nn.Linear(d_model, dim_feedforward) - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim_feedforward, d_model) - - self.norm1 = nn.LayerNorm(d_model) - self.norm2 = nn.LayerNorm(d_model) - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - - self.activation = _get_activation_fn(activation) - self.normalize_before = normalize_before - - def with_pos_embed(self, tensor, pos: Optional[Tensor]): - return tensor if pos is None else tensor + pos - - def forward_post(self, - src, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - q = k = self.with_pos_embed(src, pos) - src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, - key_padding_mask=src_key_padding_mask)[0] - src = src + self.dropout1(src2) - src = self.norm1(src) - src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) - src = src + self.dropout2(src2) - src = self.norm2(src) - return src - - def forward_pre(self, src, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - src2 = self.norm1(src) - q = k = self.with_pos_embed(src2, pos) - src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, - key_padding_mask=src_key_padding_mask)[0] - src = src + self.dropout1(src2) - src2 = self.norm2(src) - src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) - src = src + self.dropout2(src2) - return src - - def forward(self, src, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None): - if self.normalize_before: - return self.forward_pre(src, src_mask, src_key_padding_mask, pos) - return self.forward_post(src, src_mask, src_key_padding_mask, pos) - - -class TransformerDecoderLayer(nn.Module): - - def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, - activation="relu", normalize_before=False): - super().__init__() - self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) - self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) - # Implementation of Feedforward model - self.linear1 = nn.Linear(d_model, dim_feedforward) - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim_feedforward, d_model) - - self.norm1 = nn.LayerNorm(d_model) - self.norm2 = nn.LayerNorm(d_model) - self.norm3 = nn.LayerNorm(d_model) - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - self.dropout3 = nn.Dropout(dropout) - - self.activation = _get_activation_fn(activation) - self.normalize_before = normalize_before - - def with_pos_embed(self, tensor, pos: Optional[Tensor]): - return tensor if pos is None else tensor + pos - - def forward_post(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - q = k = self.with_pos_embed(tgt, query_pos) - tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, - key_padding_mask=tgt_key_padding_mask)[0] - tgt = tgt + self.dropout1(tgt2) - tgt = self.norm1(tgt) - tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), - key=self.with_pos_embed(memory, pos), - value=memory, attn_mask=memory_mask, - key_padding_mask=memory_key_padding_mask)[0] - tgt = tgt + self.dropout2(tgt2) - tgt = self.norm2(tgt) - tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) - tgt = tgt + self.dropout3(tgt2) - tgt = self.norm3(tgt) - return tgt - - def forward_pre(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - tgt2 = self.norm1(tgt) - q = k = self.with_pos_embed(tgt2, query_pos) - tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, - key_padding_mask=tgt_key_padding_mask)[0] - tgt = tgt + self.dropout1(tgt2) - tgt2 = self.norm2(tgt) - tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), - key=self.with_pos_embed(memory, pos), - value=memory, attn_mask=memory_mask, - key_padding_mask=memory_key_padding_mask)[0] - tgt = tgt + self.dropout2(tgt2) - tgt2 = self.norm3(tgt) - tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) - tgt = tgt + self.dropout3(tgt2) - return tgt - - def forward(self, tgt, memory, - tgt_mask: Optional[Tensor] = None, - memory_mask: Optional[Tensor] = None, - tgt_key_padding_mask: Optional[Tensor] = None, - memory_key_padding_mask: Optional[Tensor] = None, - pos: Optional[Tensor] = None, - query_pos: Optional[Tensor] = None): - if self.normalize_before: - return self.forward_pre(tgt, memory, tgt_mask, memory_mask, - tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) - return self.forward_post(tgt, memory, tgt_mask, memory_mask, - tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) - - -def _get_clones(module, N): - return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) - - -def build_transformer(args): - return Transformer( - d_model=args.hidden_dim, - dropout=args.dropout, - nhead=args.nheads, - dim_feedforward=args.dim_feedforward, - num_encoder_layers=args.enc_layers, - num_decoder_layers=args.dec_layers, - normalize_before=args.pre_norm, - return_intermediate_dec=True, - ) - - -def _get_activation_fn(activation): - """Return an activation function given a string""" - if activation == "relu": - return F.relu - if activation == "gelu": - return F.gelu - if activation == "glu": - return F.glu - raise RuntimeError(F"activation should be relu/gelu, not {activation}.") diff --git a/roboimi/detr/policy.py b/roboimi/detr/policy.py deleted file mode 100644 index 20ac4c0..0000000 --- a/roboimi/detr/policy.py +++ /dev/null @@ -1,163 +0,0 @@ -import torch.nn as nn -from torch.nn import functional as F -import torchvision.transforms as transforms -from torchvision.transforms import v2 -import torch -from roboimi.detr.main import build_ACT_model_and_optimizer, build_CNNMLP_model_and_optimizer - - -class ACTPolicy(nn.Module): - def __init__(self, args_override): - super().__init__() - model, optimizer = build_ACT_model_and_optimizer(args_override) - self.model = model # CVAE decoder - self.optimizer = optimizer - self.kl_weight = args_override['kl_weight'] - print(f'KL Weight {self.kl_weight}') - - def __call__(self, qpos, image, actions=None, is_pad=None): - env_state = None - normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], - std=[0.229, 0.224, 0.225]) - image = normalize(image) - if actions is not None: # training time - actions = actions[:, :self.model.num_queries] - is_pad = is_pad[:, :self.model.num_queries] - - a_hat, is_pad_hat, (mu, logvar) = self.model(qpos, image, env_state, actions, is_pad) - total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar) - loss_dict = dict() - all_l1 = F.l1_loss(actions, a_hat, reduction='none') - l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean() - loss_dict['l1'] = l1 - loss_dict['kl'] = total_kld[0] - loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight - return loss_dict - else: # inference time - a_hat, _, (_, _) = self.model(qpos, image, env_state) # no action, sample from prior - return a_hat - - def configure_optimizers(self): - return self.optimizer - -class ACTTVPolicy(nn.Module): - def __init__(self, args_override): - super().__init__() - model, optimizer = build_ACT_model_and_optimizer(args_override) - self.model = model # CVAE decoder - self.optimizer = optimizer - self.kl_weight = args_override['kl_weight'] - self.qpos_noise_std = args_override['qpos_noise_std'] - print(f'KL Weight {self.kl_weight}') - - def __call__(self, qpos, image, actions=None, is_pad=None): - env_state = None - - - - - - - - - - # normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], - # std=[0.229, 0.224, 0.225]) - # image = normalize(image) - - - patch_h = 16 - patch_w = 22 - if actions is not None: - transform = v2.Compose([ - v2.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5), - v2.RandomPerspective(distortion_scale=0.5), - v2.RandomAffine(degrees=10, translate=(0.1,0.1), scale=(0.9,1.1)), - v2.GaussianBlur(kernel_size=(9,9), sigma=(0.1,2.0)), - v2.Resize((patch_h * 14, patch_w * 14)), - # v2.CenterCrop((patch_h * 14, patch_w * 14)), - v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), - ]) - qpos += (self.qpos_noise_std**0.5)*torch.randn_like(qpos) - else: # inference time - transform = v2.Compose([ - v2.Resize((patch_h * 14, patch_w * 14)), - # v2.CenterCrop((patch_h * 14, patch_w * 14)), - v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) - ]) - - image = transform(image) - - - - - - - - - - - - - if actions is not None: # training time - actions = actions[:, :self.model.num_queries] - is_pad = is_pad[:, :self.model.num_queries] - - a_hat, is_pad_hat, (mu, logvar) = self.model(qpos, image, env_state, actions, is_pad) - total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar) - loss_dict = dict() - all_l1 = F.l1_loss(actions, a_hat, reduction='none') - l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean() - loss_dict['l1'] = l1 - loss_dict['kl'] = total_kld[0] - loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight - return loss_dict - else: # inference time - a_hat, _, (_, _) = self.model(qpos, image, env_state) # no action, sample from prior - return a_hat - - def configure_optimizers(self): - return self.optimizer - - -class CNNMLPPolicy(nn.Module): - def __init__(self, args_override): - super().__init__() - model, optimizer = build_CNNMLP_model_and_optimizer(args_override) - self.model = model # decoder - self.optimizer = optimizer - - def __call__(self, qpos, image, actions=None, is_pad=None): - env_state = None # TODO - normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], - std=[0.229, 0.224, 0.225]) - image = normalize(image) - if actions is not None: # training time - actions = actions[:, 0] - a_hat = self.model(qpos, image, env_state, actions) - mse = F.mse_loss(actions, a_hat) - loss_dict = dict() - loss_dict['mse'] = mse - loss_dict['loss'] = loss_dict['mse'] - return loss_dict - else: # inference time - a_hat = self.model(qpos, image, env_state) # no action, sample from prior - return a_hat - - def configure_optimizers(self): - return self.optimizer - -def kl_divergence(mu, logvar): - batch_size = mu.size(0) - assert batch_size != 0 - if mu.data.ndimension() == 4: - mu = mu.view(mu.size(0), mu.size(1)) - if logvar.data.ndimension() == 4: - logvar = logvar.view(logvar.size(0), logvar.size(1)) - - klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) - total_kld = klds.sum(1).mean(0, True) - dimension_wise_kld = klds.mean(0) - mean_kld = klds.mean(1).mean(0, True) - - return total_kld, dimension_wise_kld, mean_kld diff --git a/roboimi/detr/setup.py b/roboimi/detr/setup.py deleted file mode 100644 index 55d18c0..0000000 --- a/roboimi/detr/setup.py +++ /dev/null @@ -1,10 +0,0 @@ -from distutils.core import setup -from setuptools import find_packages - -setup( - name='detr', - version='0.0.0', - packages=find_packages(), - license='MIT License', - long_description=open('README.md').read(), -) \ No newline at end of file diff --git a/roboimi/detr/util/__init__.py b/roboimi/detr/util/__init__.py deleted file mode 100644 index 168f997..0000000 --- a/roboimi/detr/util/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved diff --git a/roboimi/detr/util/box_ops.py b/roboimi/detr/util/box_ops.py deleted file mode 100644 index 9c088e5..0000000 --- a/roboimi/detr/util/box_ops.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -Utilities for bounding box manipulation and GIoU. -""" -import torch -from torchvision.ops.boxes import box_area - - -def box_cxcywh_to_xyxy(x): - x_c, y_c, w, h = x.unbind(-1) - b = [(x_c - 0.5 * w), (y_c - 0.5 * h), - (x_c + 0.5 * w), (y_c + 0.5 * h)] - return torch.stack(b, dim=-1) - - -def box_xyxy_to_cxcywh(x): - x0, y0, x1, y1 = x.unbind(-1) - b = [(x0 + x1) / 2, (y0 + y1) / 2, - (x1 - x0), (y1 - y0)] - return torch.stack(b, dim=-1) - - -# modified from torchvision to also return the union -def box_iou(boxes1, boxes2): - area1 = box_area(boxes1) - area2 = box_area(boxes2) - - lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2] - rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2] - - wh = (rb - lt).clamp(min=0) # [N,M,2] - inter = wh[:, :, 0] * wh[:, :, 1] # [N,M] - - union = area1[:, None] + area2 - inter - - iou = inter / union - return iou, union - - -def generalized_box_iou(boxes1, boxes2): - """ - Generalized IoU from https://giou.stanford.edu/ - - The boxes should be in [x0, y0, x1, y1] format - - Returns a [N, M] pairwise matrix, where N = len(boxes1) - and M = len(boxes2) - """ - # degenerate boxes gives inf / nan results - # so do an early check - assert (boxes1[:, 2:] >= boxes1[:, :2]).all() - assert (boxes2[:, 2:] >= boxes2[:, :2]).all() - iou, union = box_iou(boxes1, boxes2) - - lt = torch.min(boxes1[:, None, :2], boxes2[:, :2]) - rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:]) - - wh = (rb - lt).clamp(min=0) # [N,M,2] - area = wh[:, :, 0] * wh[:, :, 1] - - return iou - (area - union) / area - - -def masks_to_boxes(masks): - """Compute the bounding boxes around the provided masks - - The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. - - Returns a [N, 4] tensors, with the boxes in xyxy format - """ - if masks.numel() == 0: - return torch.zeros((0, 4), device=masks.device) - - h, w = masks.shape[-2:] - - y = torch.arange(0, h, dtype=torch.float) - x = torch.arange(0, w, dtype=torch.float) - y, x = torch.meshgrid(y, x) - - x_mask = (masks * x.unsqueeze(0)) - x_max = x_mask.flatten(1).max(-1)[0] - x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] - - y_mask = (masks * y.unsqueeze(0)) - y_max = y_mask.flatten(1).max(-1)[0] - y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] - - return torch.stack([x_min, y_min, x_max, y_max], 1) diff --git a/roboimi/detr/util/misc.py b/roboimi/detr/util/misc.py deleted file mode 100644 index dfa9fb5..0000000 --- a/roboimi/detr/util/misc.py +++ /dev/null @@ -1,468 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -Misc functions, including distributed helpers. - -Mostly copy-paste from torchvision references. -""" -import os -import subprocess -import time -from collections import defaultdict, deque -import datetime -import pickle -from packaging import version -from typing import Optional, List - -import torch -import torch.distributed as dist -from torch import Tensor - -# needed due to empty tensor bug in pytorch and torchvision 0.5 -import torchvision -if version.parse(torchvision.__version__) < version.parse('0.7'): - from torchvision.ops import _new_empty_tensor - from torchvision.ops.misc import _output_size - - -class SmoothedValue(object): - """Track a series of values and provide access to smoothed values over a - window or the global series average. - """ - - def __init__(self, window_size=20, fmt=None): - if fmt is None: - fmt = "{median:.4f} ({global_avg:.4f})" - self.deque = deque(maxlen=window_size) - self.total = 0.0 - self.count = 0 - self.fmt = fmt - - def update(self, value, n=1): - self.deque.append(value) - self.count += n - self.total += value * n - - def synchronize_between_processes(self): - """ - Warning: does not synchronize the deque! - """ - if not is_dist_avail_and_initialized(): - return - t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda') - dist.barrier() - dist.all_reduce(t) - t = t.tolist() - self.count = int(t[0]) - self.total = t[1] - - @property - def median(self): - d = torch.tensor(list(self.deque)) - return d.median().item() - - @property - def avg(self): - d = torch.tensor(list(self.deque), dtype=torch.float32) - return d.mean().item() - - @property - def global_avg(self): - return self.total / self.count - - @property - def max(self): - return max(self.deque) - - @property - def value(self): - return self.deque[-1] - - def __str__(self): - return self.fmt.format( - median=self.median, - avg=self.avg, - global_avg=self.global_avg, - max=self.max, - value=self.value) - - -def all_gather(data): - """ - Run all_gather on arbitrary picklable data (not necessarily tensors) - Args: - data: any picklable object - Returns: - list[data]: list of data gathered from each rank - """ - world_size = get_world_size() - if world_size == 1: - return [data] - - # serialized to a Tensor - buffer = pickle.dumps(data) - storage = torch.ByteStorage.from_buffer(buffer) - tensor = torch.ByteTensor(storage).to("cuda") - - # obtain Tensor size of each rank - local_size = torch.tensor([tensor.numel()], device="cuda") - size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)] - dist.all_gather(size_list, local_size) - size_list = [int(size.item()) for size in size_list] - max_size = max(size_list) - - # receiving Tensor from all ranks - # we pad the tensor because torch all_gather does not support - # gathering tensors of different shapes - tensor_list = [] - for _ in size_list: - tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda")) - if local_size != max_size: - padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda") - tensor = torch.cat((tensor, padding), dim=0) - dist.all_gather(tensor_list, tensor) - - data_list = [] - for size, tensor in zip(size_list, tensor_list): - buffer = tensor.cpu().numpy().tobytes()[:size] - data_list.append(pickle.loads(buffer)) - - return data_list - - -def reduce_dict(input_dict, average=True): - """ - Args: - input_dict (dict): all the values will be reduced - average (bool): whether to do average or sum - Reduce the values in the dictionary from all processes so that all processes - have the averaged results. Returns a dict with the same fields as - input_dict, after reduction. - """ - world_size = get_world_size() - if world_size < 2: - return input_dict - with torch.no_grad(): - names = [] - values = [] - # sort the keys so that they are consistent across processes - for k in sorted(input_dict.keys()): - names.append(k) - values.append(input_dict[k]) - values = torch.stack(values, dim=0) - dist.all_reduce(values) - if average: - values /= world_size - reduced_dict = {k: v for k, v in zip(names, values)} - return reduced_dict - - -class MetricLogger(object): - def __init__(self, delimiter="\t"): - self.meters = defaultdict(SmoothedValue) - self.delimiter = delimiter - - def update(self, **kwargs): - for k, v in kwargs.items(): - if isinstance(v, torch.Tensor): - v = v.item() - assert isinstance(v, (float, int)) - self.meters[k].update(v) - - def __getattr__(self, attr): - if attr in self.meters: - return self.meters[attr] - if attr in self.__dict__: - return self.__dict__[attr] - raise AttributeError("'{}' object has no attribute '{}'".format( - type(self).__name__, attr)) - - def __str__(self): - loss_str = [] - for name, meter in self.meters.items(): - loss_str.append( - "{}: {}".format(name, str(meter)) - ) - return self.delimiter.join(loss_str) - - def synchronize_between_processes(self): - for meter in self.meters.values(): - meter.synchronize_between_processes() - - def add_meter(self, name, meter): - self.meters[name] = meter - - def log_every(self, iterable, print_freq, header=None): - i = 0 - if not header: - header = '' - start_time = time.time() - end = time.time() - iter_time = SmoothedValue(fmt='{avg:.4f}') - data_time = SmoothedValue(fmt='{avg:.4f}') - space_fmt = ':' + str(len(str(len(iterable)))) + 'd' - if torch.cuda.is_available(): - log_msg = self.delimiter.join([ - header, - '[{0' + space_fmt + '}/{1}]', - 'eta: {eta}', - '{meters}', - 'time: {time}', - 'data: {data}', - 'max mem: {memory:.0f}' - ]) - else: - log_msg = self.delimiter.join([ - header, - '[{0' + space_fmt + '}/{1}]', - 'eta: {eta}', - '{meters}', - 'time: {time}', - 'data: {data}' - ]) - MB = 1024.0 * 1024.0 - for obj in iterable: - data_time.update(time.time() - end) - yield obj - iter_time.update(time.time() - end) - if i % print_freq == 0 or i == len(iterable) - 1: - eta_seconds = iter_time.global_avg * (len(iterable) - i) - eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) - if torch.cuda.is_available(): - print(log_msg.format( - i, len(iterable), eta=eta_string, - meters=str(self), - time=str(iter_time), data=str(data_time), - memory=torch.cuda.max_memory_allocated() / MB)) - else: - print(log_msg.format( - i, len(iterable), eta=eta_string, - meters=str(self), - time=str(iter_time), data=str(data_time))) - i += 1 - end = time.time() - total_time = time.time() - start_time - total_time_str = str(datetime.timedelta(seconds=int(total_time))) - print('{} Total time: {} ({:.4f} s / it)'.format( - header, total_time_str, total_time / len(iterable))) - - -def get_sha(): - cwd = os.path.dirname(os.path.abspath(__file__)) - - def _run(command): - return subprocess.check_output(command, cwd=cwd).decode('ascii').strip() - sha = 'N/A' - diff = "clean" - branch = 'N/A' - try: - sha = _run(['git', 'rev-parse', 'HEAD']) - subprocess.check_output(['git', 'diff'], cwd=cwd) - diff = _run(['git', 'diff-index', 'HEAD']) - diff = "has uncommited changes" if diff else "clean" - branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD']) - except Exception: - pass - message = f"sha: {sha}, status: {diff}, branch: {branch}" - return message - - -def collate_fn(batch): - batch = list(zip(*batch)) - batch[0] = nested_tensor_from_tensor_list(batch[0]) - return tuple(batch) - - -def _max_by_axis(the_list): - # type: (List[List[int]]) -> List[int] - maxes = the_list[0] - for sublist in the_list[1:]: - for index, item in enumerate(sublist): - maxes[index] = max(maxes[index], item) - return maxes - - -class NestedTensor(object): - def __init__(self, tensors, mask: Optional[Tensor]): - self.tensors = tensors - self.mask = mask - - def to(self, device): - # type: (Device) -> NestedTensor # noqa - cast_tensor = self.tensors.to(device) - mask = self.mask - if mask is not None: - assert mask is not None - cast_mask = mask.to(device) - else: - cast_mask = None - return NestedTensor(cast_tensor, cast_mask) - - def decompose(self): - return self.tensors, self.mask - - def __repr__(self): - return str(self.tensors) - - -def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): - # TODO make this more general - if tensor_list[0].ndim == 3: - if torchvision._is_tracing(): - # nested_tensor_from_tensor_list() does not export well to ONNX - # call _onnx_nested_tensor_from_tensor_list() instead - return _onnx_nested_tensor_from_tensor_list(tensor_list) - - # TODO make it support different-sized images - max_size = _max_by_axis([list(img.shape) for img in tensor_list]) - # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) - batch_shape = [len(tensor_list)] + max_size - b, c, h, w = batch_shape - dtype = tensor_list[0].dtype - device = tensor_list[0].device - tensor = torch.zeros(batch_shape, dtype=dtype, device=device) - mask = torch.ones((b, h, w), dtype=torch.bool, device=device) - for img, pad_img, m in zip(tensor_list, tensor, mask): - pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) - m[: img.shape[1], :img.shape[2]] = False - else: - raise ValueError('not supported') - return NestedTensor(tensor, mask) - - -# _onnx_nested_tensor_from_tensor_list() is an implementation of -# nested_tensor_from_tensor_list() that is supported by ONNX tracing. -@torch.jit.unused -def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor: - max_size = [] - for i in range(tensor_list[0].dim()): - max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64) - max_size.append(max_size_i) - max_size = tuple(max_size) - - # work around for - # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) - # m[: img.shape[1], :img.shape[2]] = False - # which is not yet supported in onnx - padded_imgs = [] - padded_masks = [] - for img in tensor_list: - padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))] - padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0])) - padded_imgs.append(padded_img) - - m = torch.zeros_like(img[0], dtype=torch.int, device=img.device) - padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1) - padded_masks.append(padded_mask.to(torch.bool)) - - tensor = torch.stack(padded_imgs) - mask = torch.stack(padded_masks) - - return NestedTensor(tensor, mask=mask) - - -def setup_for_distributed(is_master): - """ - This function disables printing when not in master process - """ - import builtins as __builtin__ - builtin_print = __builtin__.print - - def print(*args, **kwargs): - force = kwargs.pop('force', False) - if is_master or force: - builtin_print(*args, **kwargs) - - __builtin__.print = print - - -def is_dist_avail_and_initialized(): - if not dist.is_available(): - return False - if not dist.is_initialized(): - return False - return True - - -def get_world_size(): - if not is_dist_avail_and_initialized(): - return 1 - return dist.get_world_size() - - -def get_rank(): - if not is_dist_avail_and_initialized(): - return 0 - return dist.get_rank() - - -def is_main_process(): - return get_rank() == 0 - - -def save_on_master(*args, **kwargs): - if is_main_process(): - torch.save(*args, **kwargs) - - -def init_distributed_mode(args): - if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: - args.rank = int(os.environ["RANK"]) - args.world_size = int(os.environ['WORLD_SIZE']) - args.gpu = int(os.environ['LOCAL_RANK']) - elif 'SLURM_PROCID' in os.environ: - args.rank = int(os.environ['SLURM_PROCID']) - args.gpu = args.rank % torch.cuda.device_count() - else: - print('Not using distributed mode') - args.distributed = False - return - - args.distributed = True - - torch.cuda.set_device(args.gpu) - args.dist_backend = 'nccl' - print('| distributed init (rank {}): {}'.format( - args.rank, args.dist_url), flush=True) - torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, - world_size=args.world_size, rank=args.rank) - torch.distributed.barrier() - setup_for_distributed(args.rank == 0) - - -@torch.no_grad() -def accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - if target.numel() == 0: - return [torch.zeros([], device=output.device)] - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0) - res.append(correct_k.mul_(100.0 / batch_size)) - return res - - -def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None): - # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor - """ - Equivalent to nn.functional.interpolate, but with support for empty batch sizes. - This will eventually be supported natively by PyTorch, and this - class can go away. - """ - if version.parse(torchvision.__version__) < version.parse('0.7'): - if input.numel() > 0: - return torch.nn.functional.interpolate( - input, size, scale_factor, mode, align_corners - ) - - output_shape = _output_size(2, input, size, scale_factor) - output_shape = list(input.shape[:-2]) + list(output_shape) - return _new_empty_tensor(input, output_shape) - else: - return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners) diff --git a/roboimi/detr/util/plot_utils.py b/roboimi/detr/util/plot_utils.py deleted file mode 100644 index 0f24bed..0000000 --- a/roboimi/detr/util/plot_utils.py +++ /dev/null @@ -1,107 +0,0 @@ -""" -Plotting utilities to visualize training logs. -""" -import torch -import pandas as pd -import numpy as np -import seaborn as sns -import matplotlib.pyplot as plt - -from pathlib import Path, PurePath - - -def plot_logs(logs, fields=('class_error', 'loss_bbox_unscaled', 'mAP'), ewm_col=0, log_name='log.txt'): - ''' - Function to plot specific fields from training log(s). Plots both training and test results. - - :: Inputs - logs = list containing Path objects, each pointing to individual dir with a log file - - fields = which results to plot from each log file - plots both training and test for each field. - - ewm_col = optional, which column to use as the exponential weighted smoothing of the plots - - log_name = optional, name of log file if different than default 'log.txt'. - - :: Outputs - matplotlib plots of results in fields, color coded for each log file. - - solid lines are training results, dashed lines are test results. - - ''' - func_name = "plot_utils.py::plot_logs" - - # verify logs is a list of Paths (list[Paths]) or single Pathlib object Path, - # convert single Path to list to avoid 'not iterable' error - - if not isinstance(logs, list): - if isinstance(logs, PurePath): - logs = [logs] - print(f"{func_name} info: logs param expects a list argument, converted to list[Path].") - else: - raise ValueError(f"{func_name} - invalid argument for logs parameter.\n \ - Expect list[Path] or single Path obj, received {type(logs)}") - - # Quality checks - verify valid dir(s), that every item in list is Path object, and that log_name exists in each dir - for i, dir in enumerate(logs): - if not isinstance(dir, PurePath): - raise ValueError(f"{func_name} - non-Path object in logs argument of {type(dir)}: \n{dir}") - if not dir.exists(): - raise ValueError(f"{func_name} - invalid directory in logs argument:\n{dir}") - # verify log_name exists - fn = Path(dir / log_name) - if not fn.exists(): - print(f"-> missing {log_name}. Have you gotten to Epoch 1 in training?") - print(f"--> full path of missing log file: {fn}") - return - - # load log file(s) and plot - dfs = [pd.read_json(Path(p) / log_name, lines=True) for p in logs] - - fig, axs = plt.subplots(ncols=len(fields), figsize=(16, 5)) - - for df, color in zip(dfs, sns.color_palette(n_colors=len(logs))): - for j, field in enumerate(fields): - if field == 'mAP': - coco_eval = pd.DataFrame( - np.stack(df.test_coco_eval_bbox.dropna().values)[:, 1] - ).ewm(com=ewm_col).mean() - axs[j].plot(coco_eval, c=color) - else: - df.interpolate().ewm(com=ewm_col).mean().plot( - y=[f'train_{field}', f'test_{field}'], - ax=axs[j], - color=[color] * 2, - style=['-', '--'] - ) - for ax, field in zip(axs, fields): - ax.legend([Path(p).name for p in logs]) - ax.set_title(field) - - -def plot_precision_recall(files, naming_scheme='iter'): - if naming_scheme == 'exp_id': - # name becomes exp_id - names = [f.parts[-3] for f in files] - elif naming_scheme == 'iter': - names = [f.stem for f in files] - else: - raise ValueError(f'not supported {naming_scheme}') - fig, axs = plt.subplots(ncols=2, figsize=(16, 5)) - for f, color, name in zip(files, sns.color_palette("Blues", n_colors=len(files)), names): - data = torch.load(f) - # precision is n_iou, n_points, n_cat, n_area, max_det - precision = data['precision'] - recall = data['params'].recThrs - scores = data['scores'] - # take precision for all classes, all areas and 100 detections - precision = precision[0, :, :, 0, -1].mean(1) - scores = scores[0, :, :, 0, -1].mean(1) - prec = precision.mean() - rec = data['recall'][0, :, 0, -1].mean() - print(f'{naming_scheme} {name}: mAP@50={prec * 100: 05.1f}, ' + - f'score={scores.mean():0.3f}, ' + - f'f1={2 * prec * rec / (prec + rec + 1e-8):0.3f}' - ) - axs[0].plot(recall, precision, c=color) - axs[1].plot(recall, scores, c=color) - - axs[0].set_title('Precision / Recall') - axs[0].legend(names) - axs[1].set_title('Scores / Recall') - axs[1].legend(names) - return fig, axs diff --git a/roboimi/envs/double_air_insert_env.py b/roboimi/envs/double_air_insert_env.py new file mode 100644 index 0000000..8896a91 --- /dev/null +++ b/roboimi/envs/double_air_insert_env.py @@ -0,0 +1,157 @@ +import copy as cp +import time + +import numpy as np + +from roboimi.envs.double_base import DualDianaMed +from roboimi.envs.double_pos_ctrl_env import DualDianaMed_Pos_Ctrl + + +SOCKET_JOINT_NAME = "blue_socket_joint" +PEG_JOINT_NAME = "red_peg_joint" +REQUIRED_TASK_STATE_KEYS = ("socket_pos", "socket_quat", "peg_pos", "peg_quat") +SOCKET_GEOM_NAMES = ("socket-1", "socket-2", "socket-3", "socket-4") +SOCKET_SUCCESS_GEOM_NAMES = ("pin",) +SOCKET_BODY_GEOM_NAMES = SOCKET_GEOM_NAMES + SOCKET_SUCCESS_GEOM_NAMES +PEG_GEOM_NAMES = ("red_peg",) +LEFT_GRIPPER_GEOM_NAMES = ( + "l_finger_left", + "r_finger_left", + "l_fingertip_g0_left", + "r_fingertip_g0_left", + "l_fingerpad_g0_left", + "r_fingerpad_g0_left", + "l_fingertip_g0_vis_left", + "r_fingertip_g0_vis_left", +) +RIGHT_GRIPPER_GEOM_NAMES = ( + "l_finger_right", + "r_finger_right", + "l_fingertip_g0_right", + "r_fingertip_g0_right", + "l_fingerpad_g0_right", + "r_fingerpad_g0_right", + "l_fingertip_g0_vis_right", + "r_fingertip_g0_vis_right", +) +TABLE_GEOM_NAME = "table" + + +def _set_free_joint_pose(joint, position, quat): + joint.qpos[:3] = np.asarray(position, dtype=np.float64) + joint.qpos[3:7] = np.asarray(quat, dtype=np.float64) + + +def set_socket_peg_task_state(mj_data, task_state): + if not isinstance(task_state, dict) or tuple(task_state.keys()) != REQUIRED_TASK_STATE_KEYS: + raise ValueError( + "task_state must be an ordered dict-like mapping with keys " + "socket_pos, socket_quat, peg_pos, peg_quat" + ) + + _set_free_joint_pose( + mj_data.joint(SOCKET_JOINT_NAME), + task_state["socket_pos"], + task_state["socket_quat"], + ) + _set_free_joint_pose( + mj_data.joint(PEG_JOINT_NAME), + task_state["peg_pos"], + task_state["peg_quat"], + ) + + +def get_socket_peg_env_state(mj_data): + socket_qpos = cp.deepcopy(np.asarray(mj_data.joint(SOCKET_JOINT_NAME).qpos[:7], dtype=np.float64)) + peg_qpos = cp.deepcopy(np.asarray(mj_data.joint(PEG_JOINT_NAME).qpos[:7], dtype=np.float64)) + return np.concatenate([socket_qpos, peg_qpos], dtype=np.float64) + + +def _normalize_contact_pairs(contact_pairs): + return {frozenset(pair) for pair in contact_pairs} + + +def _has_any_object_contact(contact_set, object_geom_names, other_geom_names): + return any( + frozenset((object_geom_name, other_geom_name)) in contact_set + for object_geom_name in object_geom_names + for other_geom_name in other_geom_names + ) + + +def _object_is_airborne(contact_set, object_geom_names): + return not _has_any_object_contact(contact_set, object_geom_names, (TABLE_GEOM_NAME,)) + + +def peg_inserted_into_socket(contact_pairs): + contact_set = _normalize_contact_pairs(contact_pairs) + return frozenset((PEG_GEOM_NAMES[0], SOCKET_SUCCESS_GEOM_NAMES[0])) in contact_set + + +def compute_air_insert_reward(contact_pairs, env_state=None): + del env_state # kept for API compatibility with rollout/eval code paths + contact_set = _normalize_contact_pairs(contact_pairs) + reward = 0 + + if _has_any_object_contact(contact_set, SOCKET_GEOM_NAMES, LEFT_GRIPPER_GEOM_NAMES): + reward += 1 + if _has_any_object_contact(contact_set, PEG_GEOM_NAMES, RIGHT_GRIPPER_GEOM_NAMES): + reward += 1 + + socket_airborne = _object_is_airborne(contact_set, SOCKET_BODY_GEOM_NAMES) + peg_airborne = _object_is_airborne(contact_set, PEG_GEOM_NAMES) + if socket_airborne: + reward += 1 + if peg_airborne: + reward += 1 + + if socket_airborne and peg_airborne and peg_inserted_into_socket(contact_pairs): + reward += 1 + + return reward + + +class DualDianaMed_Air_Insert(DualDianaMed_Pos_Ctrl): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.max_reward = 5 + + def reset(self, task_state): + set_socket_peg_task_state(self.mj_data, task_state) + DualDianaMed.reset(self) + self.top = None + self.r_vis = None + self.l_vis = None + self.front = None + if not self.is_render: + self._update_camera_images_sync() + return + self.cam_flage = True + while self.cam_flage: + if ( + type(self.top) == type(None) + or type(self.r_vis) == type(None) + or type(self.l_vis) == type(None) + or type(self.front) == type(None) + ): + time.sleep(0.001) + else: + self.cam_flage = False + + def step(self, action=np.zeros(16)): + super().step(action) + self.rew = self._get_reward() + self.obs = self._get_obs() + + def get_env_state(self): + return get_socket_peg_env_state(self.mj_data) + + def _get_reward(self): + contact_pairs = [] + for collision_num in range(self.mj_data.ncon): + geom1 = self.mj_data.contact[collision_num].geom1 + geom2 = self.mj_data.contact[collision_num].geom2 + contact_pairs.append( + (self.getID2Name("geom", geom1), self.getID2Name("geom", geom2)) + ) + return compute_air_insert_reward(contact_pairs, self.get_env_state()) diff --git a/roboimi/envs/double_base.py b/roboimi/envs/double_base.py index 1b7785b..28392b8 100644 --- a/roboimi/envs/double_base.py +++ b/roboimi/envs/double_base.py @@ -52,10 +52,11 @@ class DualDianaMed(MujocoEnv): self.r_vis = None self.l_vis = None self.top = None - self.angle = None + self.front = None self.obs = None self.rew = None + self._offscreen_renderer = None def actuate_J(self, q_target, qdot_target, Arm): @@ -160,23 +161,27 @@ class DualDianaMed(MujocoEnv): def _get_obs(self): + if not self.is_render: + self._update_camera_images_sync() obs = collections.OrderedDict() obs['qpos'] = self.get_obs_qpos obs['action'] = self.compute_qpos obs['images'] = dict() - obs['images']['top'] = self.top - obs['images']['angle'] = self.angle obs['images']['r_vis'] = self.r_vis obs['images']['l_vis'] = self.l_vis + obs['images']['top'] = self.top + obs['images']['front'] = self.front return obs def _get_image_obs(self): + if not self.is_render: + self._update_camera_images_sync() obs = collections.OrderedDict() obs['images'] = dict() - obs['images']['top'] = self.top - obs['images']['angle'] = self.angle obs['images']['r_vis'] = self.r_vis obs['images']['l_vis'] = self.l_vis + obs['images']['top'] = self.top + obs['images']['front'] = self.front return obs def _get_qpos_obs(self): @@ -194,39 +199,56 @@ class DualDianaMed(MujocoEnv): @property def cam_view(self): - if self.cam == 'top': - return self.top - elif self.cam == 'angle': - return self.angle - elif self.cam == 'r_vis': + if self.cam == 'r_vis': return self.r_vis elif self.cam == 'l_vis': return self.l_vis + elif self.cam == 'top': + return self.top + elif self.cam == 'front': + return self.front else: raise AttributeError("please input right name") + def _get_or_create_offscreen_renderer(self): + renderer = getattr(self, '_offscreen_renderer', None) + if renderer is None: + renderer = mj.Renderer(self.mj_model, height=480, width=640) + self._offscreen_renderer = renderer + return renderer + + def _render_camera_set(self, img_renderer): + img_renderer.update_scene(self.mj_data, camera="rs_cam_right") + self.r_vis = img_renderer.render()[:, :, ::-1] + img_renderer.update_scene(self.mj_data, camera="rs_cam_left") + self.l_vis = img_renderer.render()[:, :, ::-1] + img_renderer.update_scene(self.mj_data, camera="top") + self.top = img_renderer.render()[:, :, ::-1] + img_renderer.update_scene(self.mj_data, camera="front") + self.front = img_renderer.render()[:, :, ::-1] + + def _update_camera_images_sync(self): + img_renderer = self._get_or_create_offscreen_renderer() + self._render_camera_set(img_renderer) + def camera_viewer(self): - img_renderer = mj.Renderer(self.mj_model,height=480,width=640) - cv2.namedWindow('Cam view',cv2.WINDOW_NORMAL) + img_renderer = self._get_or_create_offscreen_renderer() + show_gui = self.is_render + if show_gui: + cv2.namedWindow('Cam view',cv2.WINDOW_NORMAL) while not self.exit_flag: - img_renderer.update_scene(self.mj_data,camera="rs_cam_right") - self.r_vis = img_renderer.render() - self.r_vis = self.r_vis[:, :, ::-1] - img_renderer.update_scene(self.mj_data,camera="rs_cam_left") - self.l_vis = img_renderer.render() - self.l_vis = self.l_vis[:, :, ::-1] - img_renderer.update_scene(self.mj_data,camera="top") - self.top = img_renderer.render() - self.top = self.top[:, :, ::-1] - img_renderer.update_scene(self.mj_data,camera="angle") - self.angle = img_renderer.render() - self.angle = self.angle[:, :, ::-1] - cv2.imshow('Cam view', self.cam_view) - cv2.waitKey(1) + self._render_camera_set(img_renderer) + if show_gui: + if self.cam_view is not None: + cv2.imshow('Cam view', self.cam_view) + cv2.waitKey(1) def cam_start(self): + if not self.is_render: + self.cam_thread = None + return self.cam_thread = threading.Thread(target=self.camera_viewer,daemon=True) self.cam_thread.start() @@ -291,4 +313,4 @@ if __name__ == "__main__": # print("quat_right =",quat_right,"\n") if env.is_render: env.render() - \ No newline at end of file + diff --git a/roboimi/envs/double_pos_ctrl_env.py b/roboimi/envs/double_pos_ctrl_env.py index 4d15e8c..a4fd6a4 100644 --- a/roboimi/envs/double_pos_ctrl_env.py +++ b/roboimi/envs/double_pos_ctrl_env.py @@ -34,19 +34,19 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed): is_interpolate=is_interpolate, cam_view=cam_view ) - + self.max_reward = 4 - + self.cam_start() - + def step(self,action=np.zeros(16)): action_left = self.ik_solve(action[:3],action[3:7],self.arm_left) action_right = self.ik_solve(action[7:10],action[10:14],self.arm_right) action = np.hstack((action_left,action_right,action[14:])) super().step(action) self.rew = self._get_reward() - + def step_jnt(self,action): super().step(action) @@ -63,8 +63,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed): return Arm.kdl_solver.ikSolver(p_goal, mat_goal, Arm.arm_qpos) def reset(self,box_pos): - - self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0] + + self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0] self.mj_data.joint('red_box_joint').qpos[1] = box_pos[1] self.mj_data.joint('red_box_joint').qpos[2] = box_pos[2] self.mj_data.joint('red_box_joint').qpos[3] = 1.0 @@ -72,18 +72,26 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed): self.mj_data.joint('red_box_joint').qpos[5] = 0.0 self.mj_data.joint('red_box_joint').qpos[6] = 0.0 super().reset() + self.top = None + self.r_vis = None + self.l_vis = None + self.front = None + if not self.is_render: + self._update_camera_images_sync() + return self.cam_flage = True t=0 while self.cam_flage: - if(type(self.top)==type(None) - or type(self.angle)==type(None) - or type(self.r_vis)==type(None)): + if(type(self.top)==type(None) + or type(self.r_vis)==type(None) + or type(self.l_vis)==type(None) + or type(self.front)==type(None)): time.sleep(0.001) t+=1 else: self.cam_flage=False - - + + def preStep(self, action): if isinstance(action,np.ndarray) and len(action)==16: @@ -96,7 +104,7 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed): for i in range(3): box_pose[i] = cp.deepcopy(self.mj_data.joint('red_box_joint').qpos[i]) return box_pose - + def _get_reward(self): all_contact_pairs = [] @@ -119,24 +127,36 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed): reward = 0 if touch_right_gripper and not touch_table: reward = 1 - if touch_right_gripper and not box_touch_table: + if touch_right_gripper and not box_touch_table: reward = 2 if touch_left_gripper: # attempted transfer reward = 3 if touch_left_gripper and not box_touch_table: # successful transfer reward = 4 return reward - -def make_sim_env(task_name): + +def make_sim_env(task_name, headless=False): + if task_name == 'sim_air_insert_socket_peg': + from roboimi.assets.robots.diana_med import BiDianaMedSocketPeg + from roboimi.envs.double_air_insert_env import DualDianaMed_Air_Insert + + env = DualDianaMed_Air_Insert( + robot=BiDianaMedSocketPeg(), + is_render=not headless, + control_freq=30, + is_interpolate=True, + cam_view='front' + ) + return env if 'sim_transfer' in task_name: from roboimi.assets.robots.diana_med import BiDianaMed env = DualDianaMed_Pos_Ctrl( robot=BiDianaMed(), - is_render=True, + is_render=not headless, control_freq=30, is_interpolate=True, - cam_view='angle' + cam_view='top' ) return env else: @@ -162,4 +182,3 @@ if __name__ == "__main__": env.step(action) if env.is_render: env.render() - \ No newline at end of file diff --git a/roboimi/gr00t/models/backbone.py b/roboimi/gr00t/models/backbone.py deleted file mode 100644 index 759bfb5..0000000 --- a/roboimi/gr00t/models/backbone.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -Backbone modules. -""" -from collections import OrderedDict - -import torch -import torch.nn.functional as F -import torchvision -from torch import nn -from torchvision.models._utils import IntermediateLayerGetter -from typing import Dict, List - -from util.misc import NestedTensor, is_main_process - -from .position_encoding import build_position_encoding - -class FrozenBatchNorm2d(torch.nn.Module): - """ - BatchNorm2d where the batch statistics and the affine parameters are fixed. - - Copy-paste from torchvision.misc.ops with added eps before rqsrt, - without which any other policy_models than torchvision.policy_models.resnet[18,34,50,101] - produce nans. - """ - - def __init__(self, n): - super(FrozenBatchNorm2d, self).__init__() - self.register_buffer("weight", torch.ones(n)) - self.register_buffer("bias", torch.zeros(n)) - self.register_buffer("running_mean", torch.zeros(n)) - self.register_buffer("running_var", torch.ones(n)) - - def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, - missing_keys, unexpected_keys, error_msgs): - num_batches_tracked_key = prefix + 'num_batches_tracked' - if num_batches_tracked_key in state_dict: - del state_dict[num_batches_tracked_key] - - super(FrozenBatchNorm2d, self)._load_from_state_dict( - state_dict, prefix, local_metadata, strict, - missing_keys, unexpected_keys, error_msgs) - - def forward(self, x): - # move reshapes to the beginning - # to make it fuser-friendly - w = self.weight.reshape(1, -1, 1, 1) - b = self.bias.reshape(1, -1, 1, 1) - rv = self.running_var.reshape(1, -1, 1, 1) - rm = self.running_mean.reshape(1, -1, 1, 1) - eps = 1e-5 - scale = w * (rv + eps).rsqrt() - bias = b - rm * scale - return x * scale + bias - - -class BackboneBase(nn.Module): - - def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool): - super().__init__() - # for name, parameter in backbone.named_parameters(): # only train later layers # TODO do we want this? - # if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: - # parameter.requires_grad_(False) - if return_interm_layers: - return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} - else: - return_layers = {'layer4': "0"} - self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) - self.num_channels = num_channels - - def forward(self, tensor): - xs = self.body(tensor) - return xs - # out: Dict[str, NestedTensor] = {} - # for name, x in xs.items(): - # m = tensor_list.mask - # assert m is not None - # mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] - # out[name] = NestedTensor(x, mask) - # return out - - -class Backbone(BackboneBase): - """ResNet backbone with frozen BatchNorm.""" - def __init__(self, name: str, - train_backbone: bool, - return_interm_layers: bool, - dilation: bool): - backbone = getattr(torchvision.models, name)( - replace_stride_with_dilation=[False, False, dilation], - pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d) # pretrained # TODO do we want frozen batch_norm?? - num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 - super().__init__(backbone, train_backbone, num_channels, return_interm_layers) - - -# class DINOv2BackBone(nn.Module): -# def __init__(self) -> None: -# super().__init__() -# self.body = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14') -# self.body.eval() -# self.num_channels = 384 - -# @torch.no_grad() -# def forward(self, tensor): -# xs = self.body.forward_features(tensor)["x_norm_patchtokens"] -# od = OrderedDict() -# od["0"] = xs.reshape(xs.shape[0], 22, 16, 384).permute(0, 3, 2, 1) -# return od - -class DINOv2BackBone(nn.Module): - def __init__(self, return_interm_layers: bool = False) -> None: - super().__init__() - self.body = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14') - self.body.eval() - self.num_channels = 384 - self.return_interm_layers = return_interm_layers - - @torch.no_grad() - def forward(self, tensor): - features = self.body.forward_features(tensor) - - if self.return_interm_layers: - - layer1 = features["x_norm_patchtokens"] - layer2 = features["x_norm_patchtokens"] - layer3 = features["x_norm_patchtokens"] - layer4 = features["x_norm_patchtokens"] - - od = OrderedDict() - od["0"] = layer1.reshape(layer1.shape[0], 22, 16, 384).permute(0, 3, 2, 1) - od["1"] = layer2.reshape(layer2.shape[0], 22, 16, 384).permute(0, 3, 2, 1) - od["2"] = layer3.reshape(layer3.shape[0], 22, 16, 384).permute(0, 3, 2, 1) - od["3"] = layer4.reshape(layer4.shape[0], 22, 16, 384).permute(0, 3, 2, 1) - return od - else: - xs = features["x_norm_patchtokens"] - od = OrderedDict() - od["0"] = xs.reshape(xs.shape[0], 22, 16, 384).permute(0, 3, 2, 1) - return od - -class Joiner(nn.Sequential): - def __init__(self, backbone, position_embedding): - super().__init__(backbone, position_embedding) - - def forward(self, tensor_list: NestedTensor): - xs = self[0](tensor_list) - out: List[NestedTensor] = [] - pos = [] - for name, x in xs.items(): - out.append(x) - # position encoding - pos.append(self[1](x).to(x.dtype)) - - return out, pos - - -def build_backbone(args): - position_embedding = build_position_encoding(args) - train_backbone = args.lr_backbone > 0 - return_interm_layers = args.masks - if args.backbone == 'dino_v2': - backbone = DINOv2BackBone() - else: - assert args.backbone in ['resnet18', 'resnet34'] - backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation) - model = Joiner(backbone, position_embedding) - model.num_channels = backbone.num_channels - return model diff --git a/roboimi/gr00t/models/position_encoding.py b/roboimi/gr00t/models/position_encoding.py deleted file mode 100644 index c75733e..0000000 --- a/roboimi/gr00t/models/position_encoding.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved -""" -Various positional encodings for the transformer. -""" -import math -import torch -from torch import nn - -from util.misc import NestedTensor - - -class PositionEmbeddingSine(nn.Module): - """ - This is a more standard version of the position embedding, very similar to the one - used by the Attention is all you need paper, generalized to work on images. - """ - def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): - super().__init__() - self.num_pos_feats = num_pos_feats - self.temperature = temperature - self.normalize = normalize - if scale is not None and normalize is False: - raise ValueError("normalize should be True if scale is passed") - if scale is None: - scale = 2 * math.pi - self.scale = scale - - def forward(self, tensor): - x = tensor - # mask = tensor_list.mask - # assert mask is not None - # not_mask = ~mask - - not_mask = torch.ones_like(x[0, [0]]) - y_embed = not_mask.cumsum(1, dtype=torch.float32) - x_embed = not_mask.cumsum(2, dtype=torch.float32) - if self.normalize: - eps = 1e-6 - y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale - x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale - - dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) - dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) - return pos - - -class PositionEmbeddingLearned(nn.Module): - """ - Absolute pos embedding, learned. - """ - def __init__(self, num_pos_feats=256): - super().__init__() - self.row_embed = nn.Embedding(50, num_pos_feats) - self.col_embed = nn.Embedding(50, num_pos_feats) - self.reset_parameters() - - def reset_parameters(self): - nn.init.uniform_(self.row_embed.weight) - nn.init.uniform_(self.col_embed.weight) - - def forward(self, tensor_list: NestedTensor): - x = tensor_list.tensors - h, w = x.shape[-2:] - i = torch.arange(w, device=x.device) - j = torch.arange(h, device=x.device) - x_emb = self.col_embed(i) - y_emb = self.row_embed(j) - pos = torch.cat([ - x_emb.unsqueeze(0).repeat(h, 1, 1), - y_emb.unsqueeze(1).repeat(1, w, 1), - ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1) - return pos - - -def build_position_encoding(args): - N_steps = args.hidden_dim // 2 - if args.position_embedding in ('v2', 'sine'): - # TODO find a better way of exposing other arguments - position_embedding = PositionEmbeddingSine(N_steps, normalize=True) - elif args.position_embedding in ('v3', 'learned'): - position_embedding = PositionEmbeddingLearned(N_steps) - else: - raise ValueError(f"not supported {args.position_embedding}") - - return position_embedding diff --git a/roboimi/utils/act_ex_utils.py b/roboimi/utils/act_ex_utils.py index 3c1648e..47fa832 100644 --- a/roboimi/utils/act_ex_utils.py +++ b/roboimi/utils/act_ex_utils.py @@ -1,5 +1,6 @@ import numpy as np + def sample_insertion_pose(): # Peg x_range = [0.1, 0.2] @@ -27,12 +28,31 @@ def sample_insertion_pose(): def sample_transfer_pose(): # Box - x_range = [0.0, 0.05] - y_range = [0.95, 1.05] + x_range = [-0.2, 0.2] + y_range = [0.7, 1.1] z_range = [0.47, 0.47] ranges = np.vstack([x_range, y_range, z_range]) box_position = np.random.uniform(ranges[:, 0], ranges[:, 1]) - return box_position \ No newline at end of file + return box_position + + +def sample_air_insert_socket_peg_state(): + socket_position = np.random.uniform( + low=np.array([-0.20, 0.80, 0.472], dtype=np.float32), + high=np.array([-0.10, 1.00, 0.472], dtype=np.float32), + ) + peg_position = np.random.uniform( + low=np.array([0.10, 0.80, 0.46], dtype=np.float32), + high=np.array([0.20, 1.00, 0.46], dtype=np.float32), + ) + socket_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + peg_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + return { + "socket_pos": socket_position.astype(np.float32, copy=False), + "socket_quat": socket_quat, + "peg_pos": peg_position.astype(np.float32, copy=False), + "peg_quat": peg_quat, + } diff --git a/roboimi/utils/constants.py b/roboimi/utils/constants.py index 22bc3d6..5bf065a 100644 --- a/roboimi/utils/constants.py +++ b/roboimi/utils/constants.py @@ -18,9 +18,16 @@ SIM_TASK_CONFIGS = { # }, 'sim_transfer': { 'dataset_dir': DATASET_DIR + '/sim_transfer', - 'num_episodes': 7, + 'num_episodes': 20, 'episode_len': 700, - 'camera_names': ['top','r_vis'], + 'camera_names': ['r_vis', 'top', 'front'], + 'xml_dir': HOME_PATH + '/assets' + }, + 'sim_air_insert_socket_peg': { + 'dataset_dir': DATASET_DIR + '/sim_air_insert_socket_peg', + 'num_episodes': 20, + 'episode_len': 750, + 'camera_names': ['l_vis', 'r_vis', 'front'], 'xml_dir': HOME_PATH + '/assets' }, @@ -52,13 +59,3 @@ PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x)) - -MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) -PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) - -MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE -MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)) -PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE -PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)) - -MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2 diff --git a/roboimi/utils/raw_action_trajectory_viewer.py b/roboimi/utils/raw_action_trajectory_viewer.py new file mode 100644 index 0000000..6731729 --- /dev/null +++ b/roboimi/utils/raw_action_trajectory_viewer.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +import math +import time +from pathlib import Path +from typing import Iterable + +import cv2 +import mujoco +import numpy as np + +from roboimi.assets.robots.diana_med import BiDianaMed +from roboimi.envs.mujoco_base import MujocoEnv +from roboimi.envs.double_pos_ctrl_env import make_sim_env +from roboimi.utils.act_ex_utils import sample_transfer_pose + + +def _load_raw_action_array(path: str | Path) -> np.ndarray: + path = Path(path) + if path.suffix == ".npy": + raw_action = np.load(path) + elif path.suffix == ".npz": + archive = np.load(path) + if "raw_action" in archive: + raw_action = archive["raw_action"] + elif "raw_predicted_ee_action" in archive: + raw_action = archive["raw_predicted_ee_action"] + else: + raise KeyError(f"{path} does not contain raw_action") + else: + raise ValueError(f"unsupported trajectory file: {path}") + raw_action = np.asarray(raw_action, dtype=np.float32) + if raw_action.ndim != 2 or raw_action.shape[1] < 10: + raise ValueError(f"raw_action must have shape (T, 16)-like, got {raw_action.shape}") + return raw_action + + +def disable_cv2_highgui(cv2_module=cv2): + original = { + "namedWindow": cv2_module.namedWindow, + "imshow": cv2_module.imshow, + "waitKey": cv2_module.waitKey, + } + + cv2_module.namedWindow = lambda *args, **kwargs: None + cv2_module.imshow = lambda *args, **kwargs: None + cv2_module.waitKey = lambda *args, **kwargs: 1 + + def restore(): + cv2_module.namedWindow = original["namedWindow"] + cv2_module.imshow = original["imshow"] + cv2_module.waitKey = original["waitKey"] + + return restore + + +def set_transfer_box_pose(mj_data, box_pos: np.ndarray) -> None: + box_pos = np.asarray(box_pos, dtype=np.float64) + if box_pos.shape != (3,): + raise ValueError(f"box_pos must have shape (3,), got {box_pos.shape}") + joint = mj_data.joint("red_box_joint") + joint.qpos[0] = box_pos[0] + joint.qpos[1] = box_pos[1] + joint.qpos[2] = box_pos[2] + joint.qpos[3] = 1.0 + joint.qpos[4] = 0.0 + joint.qpos[5] = 0.0 + joint.qpos[6] = 0.0 + + +def load_raw_action_positions(path: str | Path) -> dict[str, np.ndarray]: + raw_action = _load_raw_action_array(path) + return { + "left": raw_action[:, :3].astype(np.float32, copy=True), + "right": raw_action[:, 7:10].astype(np.float32, copy=True), + } + + +def _downsample_points(points: np.ndarray, stride: int) -> np.ndarray: + sampled = points[::stride] + if len(sampled) == 0: + return points + if not np.array_equal(sampled[-1], points[-1]): + sampled = np.concatenate([sampled, points[-1:]], axis=0) + return sampled + + +def build_trajectory_capsule_markers( + positions: dict[str, np.ndarray], + *, + max_markers: int, + radius: float = 0.003, + rgba: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 1.0), +) -> list[dict]: + total_segments = sum(max(len(points) - 1, 0) for points in positions.values()) + if total_segments == 0: + return [] + stride = max(1, math.ceil(total_segments / max_markers)) + markers = [] + for points in positions.values(): + sampled = _downsample_points(np.asarray(points, dtype=np.float64), stride) + for idx in range(len(sampled) - 1): + markers.append( + { + "from": sampled[idx], + "to": sampled[idx + 1], + "rgba": rgba, + "radius": float(radius), + } + ) + return markers[:max_markers] + + +def apply_capsule_markers_to_scene(user_scn, markers: Iterable[dict]) -> None: + user_scn.ngeom = 0 + for marker in markers: + if user_scn.ngeom >= user_scn.maxgeom: + break + geom = user_scn.geoms[user_scn.ngeom] + mujoco.mjv_initGeom( + geom, + mujoco.mjtGeom.mjGEOM_CAPSULE, + np.zeros(3, dtype=np.float64), + np.zeros(3, dtype=np.float64), + np.eye(3, dtype=np.float64).reshape(-1), + np.asarray(marker["rgba"], dtype=np.float32), + ) + mujoco.mjv_connector( + geom, + mujoco.mjtGeom.mjGEOM_CAPSULE, + float(marker["radius"]), + np.asarray(marker["from"], dtype=np.float64), + np.asarray(marker["to"], dtype=np.float64), + ) + user_scn.ngeom += 1 + + +def launch_raw_action_trajectory_viewer( + trajectory_path: str | Path, + *, + task_name: str = "sim_transfer", + line_radius: float = 0.004, + max_markers: int = 1500, + box_pos: np.ndarray | None = None, + disable_camera_window: bool = True, +): + positions = load_raw_action_positions(trajectory_path) + if task_name != "sim_transfer": + raise NotImplementedError(f"unsupported task_name: {task_name}") + if box_pos is None: + box_pos = sample_transfer_pose() + + robot = BiDianaMed() + viewer_env = MujocoEnv(robot=robot, is_render=True, renderer="viewer", control_freq=30) + viewer_env.reset() + set_transfer_box_pose(viewer_env.mj_data, box_pos) + mujoco.mj_forward(viewer_env.mj_model, viewer_env.mj_data) + markers = build_trajectory_capsule_markers( + positions, + max_markers=max_markers, + radius=line_radius, + ) + + if viewer_env.viewer is None or getattr(viewer_env.viewer, "user_scn", None) is None: + raise RuntimeError("viewer does not expose user_scn; cannot render trajectory overlay") + + try: + while viewer_env.viewer.is_running() and not viewer_env.exit_flag: + with viewer_env.viewer.lock(): + apply_capsule_markers_to_scene(viewer_env.viewer.user_scn, markers) + viewer_env.render() + time.sleep(1 / 60.0) + finally: + viewer_env.exit_flag = True + if getattr(viewer_env, "viewer", None) is not None: + viewer_env.viewer.close() diff --git a/roboimi/utils/streaming_episode_writer.py b/roboimi/utils/streaming_episode_writer.py new file mode 100644 index 0000000..9297069 --- /dev/null +++ b/roboimi/utils/streaming_episode_writer.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import cv2 +import h5py +import numpy as np + + +class StreamingEpisodeWriter: + """逐帧写入 episode 数据，成功后提交，失败时丢弃临时文件。""" + + def __init__( + self, + dataset_path: str | os.PathLike[str], + max_timesteps: int, + camera_names: list[str], + image_size: tuple[int, int] = (256, 256), + ) -> None: + self.dataset_path = Path(dataset_path) + self.tmp_path = Path(f"{self.dataset_path}.tmp") + self.max_timesteps = int(max_timesteps) + self.camera_names = list(camera_names) + self.image_height = int(image_size[0]) + self.image_width = int(image_size[1]) + self.frame_index = 0 + self._committed = False + self._closed = False + + self.dataset_path.parent.mkdir(parents=True, exist_ok=True) + if self.tmp_path.exists(): + self.tmp_path.unlink() + + self._file = h5py.File(self.tmp_path, "w", rdcc_nbytes=1024**2 * 2) + self._file.attrs["sim"] = True + self._file.attrs["action_repr"] = "ee_pose_xyz_quat_gripper" + self._file.attrs["image_height"] = self.image_height + self._file.attrs["image_width"] = self.image_width + self._file.attrs["camera_names"] = np.asarray(self.camera_names, dtype="S") + + observations = self._file.create_group("observations") + images = observations.create_group("images") + for cam_name in self.camera_names: + images.create_dataset( + cam_name, + (self.max_timesteps, self.image_height, self.image_width, 3), + dtype="uint8", + chunks=(1, self.image_height, self.image_width, 3), + ) + observations.create_dataset( + "qpos", + (self.max_timesteps, 16), + dtype="float32", + chunks=(min(128, self.max_timesteps), 16), + ) + self._file.create_dataset( + "action", + (self.max_timesteps, 16), + dtype="float32", + chunks=(min(128, self.max_timesteps), 16), + ) + + def append(self, qpos: np.ndarray, action: np.ndarray, images: dict[str, np.ndarray]) -> None: + if self._closed: + raise RuntimeError("writer is already closed") + if self.frame_index >= self.max_timesteps: + raise IndexError("frame index exceeds max_timesteps") + + qpos = np.asarray(qpos, dtype=np.float32) + action = np.asarray(action, dtype=np.float32) + if qpos.shape != (16,): + raise ValueError(f"qpos shape must be (16,), got {qpos.shape}") + if action.shape != (16,): + raise ValueError(f"action shape must be (16,), got {action.shape}") + + self._file["observations/qpos"][self.frame_index] = qpos + self._file["action"][self.frame_index] = action + + for cam_name in self.camera_names: + if cam_name not in images: + raise KeyError(f"missing image for camera '{cam_name}'") + self._file[f"observations/images/{cam_name}"][self.frame_index] = self._resize_image(images[cam_name]) + + self.frame_index += 1 + + def commit(self) -> None: + if self._closed: + return + self._file.flush() + self._file.close() + self._closed = True + os.replace(self.tmp_path, self.dataset_path) + self._committed = True + + def discard(self) -> None: + if not self._closed: + self._file.close() + self._closed = True + if self.tmp_path.exists(): + self.tmp_path.unlink() + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + image = np.asarray(image, dtype=np.uint8) + if image.ndim != 3 or image.shape[2] != 3: + raise ValueError(f"image shape must be HxWx3, got {image.shape}") + if image.shape[:2] == (self.image_height, self.image_width): + return image + + interpolation = cv2.INTER_AREA + if image.shape[0] < self.image_height or image.shape[1] < self.image_width: + interpolation = cv2.INTER_LINEAR + return cv2.resize(image, (self.image_width, self.image_height), interpolation=interpolation) diff --git a/roboimi/vla/__init__.py b/roboimi/vla/__init__.py new file mode 100644 index 0000000..0509741 --- /dev/null +++ b/roboimi/vla/__init__.py @@ -0,0 +1 @@ +# export VLAAgent, VLAModelConfig diff --git a/roboimi/vla/agent.py b/roboimi/vla/agent.py new file mode 100644 index 0000000..3578ac3 --- /dev/null +++ b/roboimi/vla/agent.py @@ -0,0 +1,514 @@ +import torch +import torch.nn as nn +import numpy as np +from collections import deque +from typing import Dict, Optional, Any, Tuple +from diffusers.schedulers.scheduling_ddpm import DDPMScheduler +from diffusers.schedulers.scheduling_ddim import DDIMScheduler +from roboimi.vla.models.normalization import NormalizationModule + +class VLAAgent(nn.Module): + + def __init__( + self, + vision_backbone, # 视觉编码器（ResNet 等） + state_encoder, + action_encoder, + head, + action_dim, # 机器人动作维度 (例如 7: xyz + rpy + gripper) + obs_dim, # 本体感知维度 (例如 关节角度) + pred_horizon=16, # 预测未来多少步动作 + obs_horizon=4, # 使用多少步历史观测 + diffusion_steps=100, # DDPM 加噪步数 + inference_steps=10, # DDIM 推理步数 + num_cams=3, # 视觉输入的摄像头数量 + camera_names: Optional[Tuple[str, ...]] = None, # 条件相机顺序 + dataset_stats=None, # 数据集统计信息，用于归一化 + normalization_type='min_max', # 归一化类型: 'gaussian' 或 'min_max' + num_action_steps=8, # 每次推理实际执行多少步动作 + head_type='unet', # Policy head类型: 'unet' 或 'transformer' + cond_projector=None, # 可选：将视觉+状态条件投影到head期望维度 + ): + super().__init__() + # 保存参数 + self.action_dim = action_dim + self.obs_dim = obs_dim + self.pred_horizon = pred_horizon + self.obs_horizon = obs_horizon + self.num_cams = num_cams + self.num_action_steps = num_action_steps + self.inference_steps = inference_steps + self.head_type = head_type # 'unet' 或 'transformer' + agent_camera_names = tuple(camera_names) if camera_names is not None else None + backbone_camera_names = getattr(vision_backbone, 'camera_names', None) + backbone_camera_names = tuple(backbone_camera_names) if backbone_camera_names is not None else None + backbone_num_cameras = getattr(vision_backbone, 'num_cameras', None) + if backbone_num_cameras is not None and backbone_num_cameras != self.num_cams: + raise ValueError( + f"agent.num_cams({self.num_cams}) 与 " + f"vision_backbone.num_cameras({backbone_num_cameras}) 不一致" + ) + if ( + agent_camera_names is not None + and backbone_camera_names is not None + and agent_camera_names != backbone_camera_names + ): + raise ValueError( + f"agent.camera_names({list(agent_camera_names)}) 与 " + f"vision_backbone.camera_names({list(backbone_camera_names)}) 不一致" + ) + self.camera_names = ( + agent_camera_names if agent_camera_names is not None else backbone_camera_names + ) + if self.camera_names is not None and len(self.camera_names) != self.num_cams: + raise ValueError( + f"camera_names 长度({len(self.camera_names)})与 num_cams({self.num_cams})不一致" + ) + + + # 归一化模块 - 统一训练和推理的归一化逻辑 + self.normalization = NormalizationModule( + stats=dataset_stats, + normalization_type=normalization_type + ) + + self.vision_encoder = vision_backbone + if self.camera_names is not None: + self.vision_encoder.camera_names = self.camera_names + self.condition_tokens_per_step = int(getattr(self.vision_encoder, 'tokens_per_step', 1)) + joint_vision_dim = getattr(self.vision_encoder, 'joint_output_dim', None) + if joint_vision_dim is not None: + per_token_vision_dim = int(joint_vision_dim) + self.condition_tokens_per_step = 1 + else: + single_cam_feat_dim = self.vision_encoder.output_dim + if self.condition_tokens_per_step > 1: + per_token_vision_dim = int(single_cam_feat_dim) + else: + per_token_vision_dim = int(single_cam_feat_dim) * int(num_cams) + + self.condition_sequence_length = self.obs_horizon * self.condition_tokens_per_step + self.raw_per_step_cond_dim = per_token_vision_dim + obs_dim + if cond_projector is None: + self.cond_projector = None + self.per_step_cond_dim = self.raw_per_step_cond_dim + else: + if isinstance(cond_projector, nn.Module): + self.cond_projector = cond_projector + else: + self.cond_projector = cond_projector(input_dim=self.raw_per_step_cond_dim) + self.per_step_cond_dim = self._projector_output_dim(self.cond_projector, self.raw_per_step_cond_dim) + + # global_cond_dim: 展平后的总维度（用于UNet） + self.global_cond_dim = self.per_step_cond_dim * self.condition_sequence_length + + self.noise_scheduler = DDPMScheduler( + num_train_timesteps=diffusion_steps, + beta_schedule='squaredcos_cap_v2', # 机器人任务常用的 schedule + clip_sample=True, + prediction_type='epsilon' # 预测噪声 + ) + + # DDIM 调度器用于快速推理 + self.infer_scheduler = DDIMScheduler( + num_train_timesteps=diffusion_steps, + beta_schedule='squaredcos_cap_v2', + clip_sample=True, + prediction_type='epsilon' + ) + + # 根据head类型初始化不同的参数 + if head_type == 'transformer': + # 如果head已经是nn.Module实例，直接使用；否则需要初始化 + if isinstance(head, nn.Module): + # 已经是实例化的模块（测试时直接传入�� + self.noise_pred_net = head + else: + # Hydra部分初始化的对象，调用时传入参数 + self.noise_pred_net = head( + input_dim=action_dim, + output_dim=action_dim, + horizon=pred_horizon, + n_obs_steps=self.condition_sequence_length, + cond_dim=self.per_step_cond_dim # 每步的条件维度 + ) + else: # 'unet' (default) + # UNet接口: input_dim, global_cond_dim + self.noise_pred_net = head( + input_dim=action_dim, + global_cond_dim=self.global_cond_dim + ) + + self.state_encoder = state_encoder + self.action_encoder = action_encoder + + # 初始化队列（用于在线推理） + self.reset() + + def _get_model_device(self) -> torch.device: + """获取模型当前所在设备。""" + return next(self.parameters()).device + + def _move_to_device(self, data, device: torch.device): + """递归地将张量数据移动到指定设备。""" + if torch.is_tensor(data): + return data.to(device) + if isinstance(data, dict): + return {k: self._move_to_device(v, device) for k, v in data.items()} + if isinstance(data, list): + return [self._move_to_device(v, device) for v in data] + if isinstance(data, tuple): + return tuple(self._move_to_device(v, device) for v in data) + return data + + @staticmethod + def _projector_output_dim(projector: nn.Module, fallback: int) -> int: + output_dim = getattr(projector, 'output_dim', None) + if output_dim is not None: + return int(output_dim) + out_features = getattr(projector, 'out_features', None) + if out_features is not None: + return int(out_features) + linear = getattr(projector, 'linear', None) + linear_out_features = getattr(linear, 'out_features', None) + if linear_out_features is not None: + return int(linear_out_features) + return int(fallback) + + def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """按显式配置的相机顺序返回图像字典。""" + if self.camera_names is None: + camera_names = tuple(sorted(images.keys())) + if len(camera_names) != self.num_cams: + raise ValueError( + f"图像条件相机数量({len(camera_names)})与 num_cams({self.num_cams})不一致" + ) + return {cam_name: images[cam_name] for cam_name in camera_names} + + missing = [cam_name for cam_name in self.camera_names if cam_name not in images] + if missing: + raise ValueError( + f"图像条件缺少必需相机。missing={missing}, expected={list(self.camera_names)}" + ) + return {cam_name: images[cam_name] for cam_name in self.camera_names} + + def _build_cond(self, images: Dict[str, torch.Tensor], states: torch.Tensor) -> torch.Tensor: + """构造每步条件，确保图像条件顺序稳定。""" + ordered_images = self._order_images(images) + visual_features = self.vision_encoder(ordered_images) + state_features = self.state_encoder(states) + if visual_features.ndim == 4: + batch_size, obs_steps, token_count, _ = visual_features.shape + if obs_steps != state_features.shape[1]: + raise RuntimeError( + f"观测时间维不匹配: visual={obs_steps}, state={state_features.shape[1]}" + ) + if token_count != self.condition_tokens_per_step: + raise RuntimeError( + f"条件token数量不匹配: got {token_count}, expected {self.condition_tokens_per_step}" + ) + state_features = state_features.unsqueeze(2).expand(-1, -1, token_count, -1) + cond = torch.cat([visual_features, state_features], dim=-1) + if cond.shape[-1] != self.raw_per_step_cond_dim: + raise RuntimeError( + f"原始条件维度不匹配: got {cond.shape[-1]}, expected {self.raw_per_step_cond_dim}" + ) + if self.cond_projector is not None: + cond = self.cond_projector(cond) + if cond.shape[-1] != self.per_step_cond_dim: + raise RuntimeError( + f"条件维度不匹配: got {cond.shape[-1]}, expected {self.per_step_cond_dim}" + ) + cond = cond.reshape(batch_size, obs_steps * token_count, self.per_step_cond_dim) + expected_length = self.condition_sequence_length + if cond.shape[1] != expected_length: + raise RuntimeError( + f"条件序列长度不匹配: got {cond.shape[1]}, expected {expected_length}" + ) + return cond + + cond = torch.cat([visual_features, state_features], dim=-1) + if cond.shape[-1] != self.raw_per_step_cond_dim: + raise RuntimeError( + f"原始条件维度不匹配: got {cond.shape[-1]}, expected {self.raw_per_step_cond_dim}" + ) + if self.cond_projector is not None: + cond = self.cond_projector(cond) + if cond.shape[-1] != self.per_step_cond_dim: + raise RuntimeError( + f"条件维度不匹配: got {cond.shape[-1]}, expected {self.per_step_cond_dim}" + ) + return cond + + # ========================== + # 训练阶段 (Training) + # ========================== + def compute_loss(self, batch): + """ + 计算训练损失 + + Args: + batch: 包含 images, qpos (本体感知), action, action_is_pad 的字典 + """ + actions, states, images = batch['action'], batch['qpos'], batch['images'] + action_is_pad = batch.get('action_is_pad', None) # 获取padding mask + B = actions.shape[0] + + # 归一化 states (qpos) 和 actions + states = self.normalization.normalize_qpos(states) + actions = self.normalization.normalize_action(actions) + + # 1. 提取视觉特征 + per_step_cond = self._build_cond(images, states) + action_features = self.action_encoder(actions) + + # 2. 采样噪声 + noise = torch.randn_like(action_features) + + # 3. 随机采样时间步 (Timesteps) + timesteps = torch.randint( + 0, self.noise_scheduler.config.num_train_timesteps, + (B,), device=action_features.device + ).long() + + # 4. 给动作加噪 (Forward Diffusion) + noisy_actions = self.noise_scheduler.add_noise( + action_features, noise, timesteps + ) + + # 拼接全局条件并展平 + # per_step_cond: (B, obs_horizon, vision_dim * num_cams + obs_dim) + # 展平后用于 UNet，全序列形式用于 Transformer + global_cond = per_step_cond.flatten(start_dim=1) + + # 5. 网络预测噪声（根据head类型选择接口） + if self.head_type == 'transformer': + pred_noise = self.noise_pred_net( + sample=noisy_actions, + timestep=timesteps, + cond=per_step_cond + ) + else: # 'unet' + pred_noise = self.noise_pred_net( + sample=noisy_actions, + timestep=timesteps, + global_cond=global_cond + ) + + # 6. 计算 Loss (MSE)，支持 padding mask + loss = nn.functional.mse_loss(pred_noise, noise, reduction='none') + + # 如果提供了 action_is_pad，对padding位置进行mask + if action_is_pad is not None: + # action_is_pad: (B, pred_horizon)，扩展到 (B, pred_horizon, action_dim) + mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype) # 1.0表示有效数据 + valid_count = mask.sum() * loss.shape[-1] + loss = (loss * mask).sum() / valid_count.clamp_min(1.0) + else: + loss = loss.mean() + + return loss + + # ========================== + # 队列管理 (Queue Management) + # ========================== + def reset(self): + """清空观测和动作队列。应在 env.reset() 时调用""" + self._queues = { + 'qpos': deque(maxlen=self.obs_horizon), + 'images': deque(maxlen=self.obs_horizon), + 'action': deque(maxlen=self.pred_horizon - self.obs_horizon + 1), # 可执行的动作缓存 + } + + def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None: + """ + 将新的观测添加到队列中。 + + Args: + observation: 包含 'qpos' 和 'images' 的字典 + """ + # 添加本体感知 + if 'qpos' in observation: + self._queues['qpos'].append(observation['qpos'].clone()) + + # 添加图像 + if 'images' in observation: + ordered_images = self._order_images(observation['images']) + self._queues['images'].append({k: v.clone() for k, v in ordered_images.items()}) + + def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]: + """ + 从队列中准备用于推理的批量观测。 + 如果队列未满（首次调用时），用最新观测重复填充。 + + Returns: + batch: 包含堆叠后的历史观测的字典 + """ + # 堆叠历史本体感知 + qpos_list = list(self._queues['qpos']) + if len(qpos_list) == 0: + raise ValueError("观测队列为空，请先调用 _populate_queues 添加观测") + # 如果队列未满，用最后一个观测填充 + while len(qpos_list) < self.obs_horizon: + qpos_list.append(qpos_list[-1]) + batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0) # (1, obs_horizon, obs_dim) + + # 堆叠历史图像 + images_list = list(self._queues['images']) + if len(images_list) == 0: + raise ValueError("图像队列为空，请先调用 _populate_queues 添加观测") + # 如果队列未满，用最后一个观测填充 + while len(images_list) < self.obs_horizon: + images_list.append(images_list[-1]) + + batch_images = {} + camera_names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys())) + for cam_name in camera_names: + batch_images[cam_name] = torch.stack([img[cam_name] for img in images_list], dim=0).unsqueeze(0) + + return {'qpos': batch_qpos, 'images': batch_images} + + # ========================== + # 在线推理 (Online Inference) + # ========================== + @torch.no_grad() + def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor: + """ + 根据当前观测选择单个动作。 + + 这个方法维护一个历史观测和生成动作轨迹的缓存。工作流程： + - 缓存 `obs_horizon` 步的历史观测 + - Diffusion 模型生成 `pred_horizon` 步的动作 + - 实际执行 `num_action_steps` 步动作 + + 示意图: + -------------------------------------------------------------- + (图例: o=obs_horizon, h=pred_horizon, a=num_action_steps) + |时间步 | 0 | 1 | ... | o-1 | o | ... | h-1 | + |观测是否使用 | 是 | 是 | 是 | 是 | 否 | 否 | 否 | + |动作是否生成 | 是 | 是 | 是 | 是 | 是 | 是 | 是 | + |动作是否执行 | 否 | 否 | 否 | 否 | 是 | 是 | 是 | + -------------------------------------------------------------- + + Args: + observation: 包含 'qpos' 和 'images' 的字典 + + Returns: + action: (action_dim,) 单个动作 + """ + # 使用模型当前设备作为唯一真值，将输入移动到模型设备 + # 避免根据CPU观测把模型错误搬回CPU。 + device = self._get_model_device() + observation = self._move_to_device(observation, device) + + # 将新观测添加到队列 + self._populate_queues(observation) + + # 如果动作队列为空，生成新的动作序列 + if len(self._queues['action']) == 0: + # 从队列准备批量观测 + batch = self._prepare_observation_batch() + + # 生成动作块 + actions = self.predict_action_chunk(batch) # (1, pred_horizon, action_dim) + + # 提取可执行的动作部分 + # 从 obs_horizon-1 开始，因为前面的动作对应过去的观测 + start = self.obs_horizon - 1 + end = start + self.num_action_steps + executable_actions = actions[:, start:end] # (1, num_action_steps, action_dim) + + # 将动作添加到队列 + for i in range(executable_actions.shape[1]): + self._queues['action'].append(executable_actions[:, i].squeeze(0)) # (action_dim,) + + # 从队列中取出一个动作 + action = self._queues['action'].popleft() # (action_dim,) + + return action + + @torch.no_grad() + def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + """ + 预测一个动作块（用于在线推理）。 + + Args: + batch: 包含 'qpos' 和 'images' 的字典 + - qpos: (B, obs_horizon, obs_dim) + - images: Dict[str, (B, obs_horizon, C, H, W)] + + Returns: + actions: (B, pred_horizon, action_dim) 预测的动作序列 + """ + return self.predict_action(batch['images'], batch['qpos']) + + # ========================== + # 批量推理 (Batch Inference - 原有方法) + # ========================== + @torch.no_grad() + def predict_action(self, images, proprioception): + """ + 批量预测动作序列（用于训练和离线评估） + + Args: + images: 图像观测字典 + proprioception: 本体感知观测 (qpos) + + Returns: + denormalized_actions: 反归一化后的动作序列 + """ + B = proprioception.shape[0] + + # 归一化 proprioception (qpos) + proprioception = self.normalization.normalize_qpos(proprioception) + + # 1. 提取当前观测特征（只提取一次） + per_step_cond = self._build_cond(images, proprioception) + + # 拼接条件（只计算一次） + global_cond_flat = per_step_cond.flatten(start_dim=1) + if self.head_type == 'transformer': + cond = per_step_cond + else: + cond = None + + # 2. 初始化纯高斯噪声动作 + # 形状: (B, pred_horizon, action_dim) + device = per_step_cond.device + current_actions = torch.randn( + (B, self.pred_horizon, self.action_dim), device=device + ) + + # 3. 逐步去噪循环 (Reverse Diffusion) + self.infer_scheduler.set_timesteps(self.inference_steps) # DDIM 推理步数 + + for t in self.infer_scheduler.timesteps: + model_input = current_actions + + # 预测噪声（根据head类型选择接口） + if self.head_type == 'transformer': + noise_pred = self.noise_pred_net( + sample=model_input, + timestep=t, + cond=cond + ) + else: # 'unet' + noise_pred = self.noise_pred_net( + sample=model_input, + timestep=t, + global_cond=global_cond_flat + ) + + # 移除噪声，更新 current_actions + current_actions = self.infer_scheduler.step( + noise_pred, t, current_actions + ).prev_sample + + # 4. 反归一化动作序列 + denormalized_actions = self.normalization.denormalize_action(current_actions) + + return denormalized_actions + + def get_normalization_stats(self): + """获取归一化统计信息（用于保存到 checkpoint）""" + return self.normalization.get_stats() diff --git a/roboimi/vla/agent_gr00t_dit.py b/roboimi/vla/agent_gr00t_dit.py new file mode 100644 index 0000000..eadfad8 --- /dev/null +++ b/roboimi/vla/agent_gr00t_dit.py @@ -0,0 +1,217 @@ +import torch +import torch.nn as nn +from collections import deque +from typing import Dict + +from diffusers.schedulers.scheduling_ddpm import DDPMScheduler +from diffusers.schedulers.scheduling_ddim import DDIMScheduler + +from roboimi.vla.models.normalization import NormalizationModule + + +class VLAAgentGr00tDiT(nn.Module): + """ + VLA Agent variant that swaps Transformer1D head with gr00t DiT head. + Other components (backbone/encoders/scheduler/queue logic) stay aligned + with the existing VLAAgent implementation. + """ + + def __init__( + self, + vision_backbone, + state_encoder, + action_encoder, + head, + action_dim, + obs_dim, + pred_horizon=16, + obs_horizon=4, + diffusion_steps=100, + inference_steps=10, + num_cams=3, + dataset_stats=None, + normalization_type="min_max", + num_action_steps=8, + ): + super().__init__() + self.action_dim = action_dim + self.obs_dim = obs_dim + self.pred_horizon = pred_horizon + self.obs_horizon = obs_horizon + self.num_cams = num_cams + self.num_action_steps = num_action_steps + self.inference_steps = inference_steps + + self.normalization = NormalizationModule( + stats=dataset_stats, + normalization_type=normalization_type, + ) + + self.vision_encoder = vision_backbone + single_cam_feat_dim = self.vision_encoder.output_dim + self.per_step_cond_dim = single_cam_feat_dim * num_cams + obs_dim + + self.noise_scheduler = DDPMScheduler( + num_train_timesteps=diffusion_steps, + beta_schedule="squaredcos_cap_v2", + clip_sample=True, + prediction_type="epsilon", + ) + self.infer_scheduler = DDIMScheduler( + num_train_timesteps=diffusion_steps, + beta_schedule="squaredcos_cap_v2", + clip_sample=True, + prediction_type="epsilon", + ) + + if isinstance(head, nn.Module): + self.noise_pred_net = head + else: + self.noise_pred_net = head( + input_dim=action_dim, + output_dim=action_dim, + horizon=pred_horizon, + n_obs_steps=obs_horizon, + cond_dim=self.per_step_cond_dim, + ) + + self.state_encoder = state_encoder + self.action_encoder = action_encoder + self.reset() + + def _get_model_device(self) -> torch.device: + return next(self.parameters()).device + + def _move_to_device(self, data, device: torch.device): + if torch.is_tensor(data): + return data.to(device) + if isinstance(data, dict): + return {k: self._move_to_device(v, device) for k, v in data.items()} + if isinstance(data, list): + return [self._move_to_device(v, device) for v in data] + if isinstance(data, tuple): + return tuple(self._move_to_device(v, device) for v in data) + return data + + def _build_cond(self, images: Dict[str, torch.Tensor], states: torch.Tensor) -> torch.Tensor: + visual_features = self.vision_encoder(images) + state_features = self.state_encoder(states) + return torch.cat([visual_features, state_features], dim=-1) + + def compute_loss(self, batch): + actions, states, images = batch["action"], batch["qpos"], batch["images"] + action_is_pad = batch.get("action_is_pad", None) + bsz = actions.shape[0] + + states = self.normalization.normalize_qpos(states) + actions = self.normalization.normalize_action(actions) + + action_features = self.action_encoder(actions) + cond = self._build_cond(images, states) + + noise = torch.randn_like(action_features) + timesteps = torch.randint( + 0, + self.noise_scheduler.config.num_train_timesteps, + (bsz,), + device=action_features.device, + ).long() + noisy_actions = self.noise_scheduler.add_noise(action_features, noise, timesteps) + + pred_noise = self.noise_pred_net( + sample=noisy_actions, + timestep=timesteps, + cond=cond, + ) + loss = nn.functional.mse_loss(pred_noise, noise, reduction="none") + + if action_is_pad is not None: + mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype) + valid_count = mask.sum() * loss.shape[-1] + loss = (loss * mask).sum() / valid_count.clamp_min(1.0) + else: + loss = loss.mean() + + return loss + + def reset(self): + self._queues = { + "qpos": deque(maxlen=self.obs_horizon), + "images": deque(maxlen=self.obs_horizon), + "action": deque(maxlen=self.pred_horizon - self.obs_horizon + 1), + } + + def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None: + if "qpos" in observation: + self._queues["qpos"].append(observation["qpos"].clone()) + if "images" in observation: + self._queues["images"].append({k: v.clone() for k, v in observation["images"].items()}) + + def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]: + qpos_list = list(self._queues["qpos"]) + if len(qpos_list) == 0: + raise ValueError("observation queue is empty.") + while len(qpos_list) < self.obs_horizon: + qpos_list.append(qpos_list[-1]) + batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0) + + images_list = list(self._queues["images"]) + if len(images_list) == 0: + raise ValueError("image queue is empty.") + while len(images_list) < self.obs_horizon: + images_list.append(images_list[-1]) + + batch_images = {} + for cam_name in images_list[0].keys(): + batch_images[cam_name] = torch.stack( + [img[cam_name] for img in images_list], dim=0 + ).unsqueeze(0) + + return {"qpos": batch_qpos, "images": batch_images} + + @torch.no_grad() + def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor: + device = self._get_model_device() + observation = self._move_to_device(observation, device) + self._populate_queues(observation) + + if len(self._queues["action"]) == 0: + batch = self._prepare_observation_batch() + actions = self.predict_action_chunk(batch) + start = self.obs_horizon - 1 + end = start + self.num_action_steps + executable_actions = actions[:, start:end] + for i in range(executable_actions.shape[1]): + self._queues["action"].append(executable_actions[:, i].squeeze(0)) + + return self._queues["action"].popleft() + + @torch.no_grad() + def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + return self.predict_action(batch["images"], batch["qpos"]) + + @torch.no_grad() + def predict_action(self, images, proprioception): + bsz = proprioception.shape[0] + proprioception = self.normalization.normalize_qpos(proprioception) + cond = self._build_cond(images, proprioception) + + device = cond.device + current_actions = torch.randn((bsz, self.pred_horizon, self.action_dim), device=device) + self.infer_scheduler.set_timesteps(self.inference_steps) + + for t in self.infer_scheduler.timesteps: + noise_pred = self.noise_pred_net( + sample=current_actions, + timestep=t, + cond=cond, + ) + current_actions = self.infer_scheduler.step( + noise_pred, t, current_actions + ).prev_sample + + return self.normalization.denormalize_action(current_actions) + + def get_normalization_stats(self): + return self.normalization.get_stats() + diff --git a/roboimi/vla/agent_imf.py b/roboimi/vla/agent_imf.py new file mode 100644 index 0000000..6dfc307 --- /dev/null +++ b/roboimi/vla/agent_imf.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from contextlib import nullcontext +from typing import Dict, Optional + +import torch +import torch.nn.functional as F + +from roboimi.vla.agent import VLAAgent + +try: + from torch.func import jvp as TORCH_FUNC_JVP +except ImportError: # pragma: no cover + TORCH_FUNC_JVP = None + + +class IMFVLAAgent(VLAAgent): + def __init__(self, *args, inference_steps: int = 1, **kwargs): + if inference_steps != 1: + raise ValueError( + 'IMFVLAAgent only supports one-step inference; ' + f'inference_steps must be 1, got {inference_steps}.' + ) + super().__init__(*args, inference_steps=inference_steps, **kwargs) + self.inference_steps = 1 + + @staticmethod + def _broadcast_batch_time(value: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + while value.ndim < reference.ndim: + value = value.unsqueeze(-1) + return value + + @staticmethod + def _apply_conditioning( + trajectory: torch.Tensor, + condition_data: Optional[torch.Tensor] = None, + condition_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if condition_data is None or condition_mask is None: + return trajectory + conditioned = trajectory.clone() + conditioned[condition_mask] = condition_data[condition_mask] + return conditioned + + @staticmethod + def _jvp_math_sdp_context(z_t: torch.Tensor): + if z_t.is_cuda: + return torch.backends.cuda.sdp_kernel( + enable_flash=False, + enable_math=True, + enable_mem_efficient=False, + enable_cudnn=False, + ) + return nullcontext() + + @staticmethod + def _jvp_tangents(v: torch.Tensor, r: torch.Tensor, t: torch.Tensor): + return v.detach(), torch.zeros_like(r), torch.ones_like(t) + + def fn(self, z: torch.Tensor, r: torch.Tensor, t: torch.Tensor, cond=None) -> torch.Tensor: + return self.noise_pred_net(z, r, t, cond=cond) + + def _compute_u_and_du_dt( + self, + z_t: torch.Tensor, + r: torch.Tensor, + t: torch.Tensor, + cond, + v: torch.Tensor, + condition_data: Optional[torch.Tensor] = None, + condition_mask: Optional[torch.Tensor] = None, + ): + tangents = self._jvp_tangents(v, r, t) + + def g(z, r_value, t_value): + conditioned_z = self._apply_conditioning(z, condition_data, condition_mask) + return self.fn(conditioned_z, r_value, t_value, cond=cond) + + with self._jvp_math_sdp_context(z_t): + if TORCH_FUNC_JVP is not None: + try: + return TORCH_FUNC_JVP(g, (z_t, r, t), tangents) + except (RuntimeError, TypeError, NotImplementedError): + pass + + u = g(z_t, r, t) + _, du_dt = torch.autograd.functional.jvp( + g, + (z_t, r, t), + tangents, + create_graph=False, + strict=False, + ) + return u, du_dt + + def _compound_velocity( + self, + u: torch.Tensor, + du_dt: torch.Tensor, + r: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + delta = self._broadcast_batch_time(t - r, u) + return u + delta * du_dt.detach() + + def _sample_one_step( + self, + z_t: torch.Tensor, + r: Optional[torch.Tensor] = None, + t: Optional[torch.Tensor] = None, + cond=None, + ) -> torch.Tensor: + batch_size = z_t.shape[0] + if t is None: + t = torch.ones(batch_size, device=z_t.device, dtype=z_t.dtype) + if r is None: + r = torch.zeros(batch_size, device=z_t.device, dtype=z_t.dtype) + u = self.fn(z_t, r, t, cond=cond) + delta = self._broadcast_batch_time(t - r, z_t) + return z_t - delta * u + + def compute_loss(self, batch): + actions, states, images = batch['action'], batch['qpos'], batch['images'] + action_is_pad = batch.get('action_is_pad', None) + batch_size = actions.shape[0] + + states = self.normalization.normalize_qpos(states) + actions = self.normalization.normalize_action(actions) + cond = self._build_cond(images, states) + + x = actions + e = torch.randn_like(x) + t = torch.rand(batch_size, device=x.device, dtype=x.dtype) + r = torch.rand(batch_size, device=x.device, dtype=x.dtype) + t, r = torch.maximum(t, r), torch.minimum(t, r) + + t_broadcast = self._broadcast_batch_time(t, x) + z_t = (1 - t_broadcast) * x + t_broadcast * e + + v = self.fn(z_t, t, t, cond=cond) + u, du_dt = self._compute_u_and_du_dt(z_t, r, t, cond=cond, v=v) + V = self._compound_velocity(u, du_dt, r, t) + target = e - x + + loss = F.mse_loss(V, target, reduction='none') + if action_is_pad is not None: + mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype) + valid_count = mask.sum() * loss.shape[-1] + loss = (loss * mask).sum() / valid_count.clamp_min(1.0) + else: + loss = loss.mean() + return loss + + @torch.no_grad() + def predict_action(self, images, proprioception): + batch_size = proprioception.shape[0] + proprioception = self.normalization.normalize_qpos(proprioception) + cond = self._build_cond(images, proprioception) + z_t = torch.randn((batch_size, self.pred_horizon, self.action_dim), device=cond.device, dtype=cond.dtype) + action = self._sample_one_step(z_t, cond=cond) + return self.normalization.denormalize_action(action) diff --git a/roboimi/vla/conf/agent/lewm_imf_attnres.yaml b/roboimi/vla/conf/agent/lewm_imf_attnres.yaml new file mode 100644 index 0000000..f7394e8 --- /dev/null +++ b/roboimi/vla/conf/agent/lewm_imf_attnres.yaml @@ -0,0 +1,41 @@ +# @package agent +defaults: + - /backbone@vision_backbone: lewm_vit_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /head: imf_transformer1d + - _self_ + +_target_: roboimi.vla.agent_imf.IMFVLAAgent + +action_dim: 16 +obs_dim: 16 +normalization_type: "min_max" +pred_horizon: 16 +obs_horizon: 2 +num_action_steps: 8 +camera_names: ${data.camera_names} +num_cams: 3 + +vision_backbone: + num_cameras: ${agent.num_cams} + camera_names: ${agent.camera_names} + fused_camera_names: [front, top, r_vis] + +diffusion_steps: 100 +inference_steps: 1 +head_type: "transformer" + +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + n_obs_steps: ${agent.obs_horizon} + cond_dim: 208 + causal_attn: false + time_as_cond: true + obs_as_cond: true + n_cond_layers: 0 + backbone_type: attnres_full + n_head: 1 + n_kv_head: 1 diff --git a/roboimi/vla/conf/agent/resnet_diffusion.yaml b/roboimi/vla/conf/agent/resnet_diffusion.yaml new file mode 100644 index 0000000..3574f96 --- /dev/null +++ b/roboimi/vla/conf/agent/resnet_diffusion.yaml @@ -0,0 +1,39 @@ +# @package agent +defaults: + # - /backbone@vision_backbone: resnet + - /backbone@vision_backbone: resnet_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /head: conditional_unet1d + - _self_ + +_target_: roboimi.vla.agent.VLAAgent + +# ==================== +# 模型维度配置 +# ==================== +action_dim: 16 # 动作维度（机器人关节数） +obs_dim: 16 # 本体感知维度（关节位置） + +# ==================== +# +# ==================== +normalization_type: "min_max" # "min_max" or "gaussian" + +# ==================== +# 时间步配置 +# ==================== +pred_horizon: 16 # 预测未来多少步动作 +obs_horizon: 2 # 使用多少步历史观测 +num_action_steps: 8 # 每次推理实际执行多少步动作（应 <= pred_horizon - obs_horizon + 1） + +# ==================== +# 相机配置 +# ==================== +num_cams: 3 # 摄像头数量 (r_vis, top, front) + +# ==================== +# 扩散过程配置 +# ==================== +diffusion_steps: 100 # 扩散训练步数（DDPM） +inference_steps: 10 # 推理时的去噪步数（DDIM，固定为 10） \ No newline at end of file diff --git a/roboimi/vla/conf/agent/resnet_gr00t_dit.yaml b/roboimi/vla/conf/agent/resnet_gr00t_dit.yaml new file mode 100644 index 0000000..e21f39f --- /dev/null +++ b/roboimi/vla/conf/agent/resnet_gr00t_dit.yaml @@ -0,0 +1,37 @@ +# @package agent +defaults: + - /backbone@vision_backbone: resnet_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /head: gr00t_dit1d + - _self_ + +_target_: roboimi.vla.agent_gr00t_dit.VLAAgentGr00tDiT + +# Model dimensions +action_dim: 16 +obs_dim: 16 + +# Normalization +normalization_type: "min_max" + +# Horizons +pred_horizon: 16 +obs_horizon: 2 +num_action_steps: 8 + +# Cameras +num_cams: 3 + +# Diffusion +diffusion_steps: 100 +inference_steps: 10 + +# Head overrides +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + n_obs_steps: ${agent.obs_horizon} + cond_dim: 208 + diff --git a/roboimi/vla/conf/agent/resnet_imf_attnres.yaml b/roboimi/vla/conf/agent/resnet_imf_attnres.yaml new file mode 100644 index 0000000..e04bfb4 --- /dev/null +++ b/roboimi/vla/conf/agent/resnet_imf_attnres.yaml @@ -0,0 +1,40 @@ +# @package agent +defaults: + - /backbone@vision_backbone: resnet_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /head: imf_transformer1d + - _self_ + +_target_: roboimi.vla.agent_imf.IMFVLAAgent + +action_dim: 16 +obs_dim: 16 +normalization_type: "min_max" +pred_horizon: 16 +obs_horizon: 2 +num_action_steps: 8 +camera_names: ${data.camera_names} +num_cams: 3 + +vision_backbone: + num_cameras: ${agent.num_cams} + camera_names: ${agent.camera_names} + +diffusion_steps: 100 +inference_steps: 1 +head_type: "transformer" + +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + n_obs_steps: ${agent.obs_horizon} + cond_dim: 208 + causal_attn: false + time_as_cond: true + obs_as_cond: true + n_cond_layers: 0 + backbone_type: attnres_full + n_head: 1 + n_kv_head: 1 diff --git a/roboimi/vla/conf/agent/resnet_imf_attnres_multitoken.yaml b/roboimi/vla/conf/agent/resnet_imf_attnres_multitoken.yaml new file mode 100644 index 0000000..bfb5906 --- /dev/null +++ b/roboimi/vla/conf/agent/resnet_imf_attnres_multitoken.yaml @@ -0,0 +1,48 @@ +# @package agent +defaults: + - /backbone@vision_backbone: resnet_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /modules@cond_projector: linear_condition_projector + - /head: imf_transformer1d + - _self_ + +_target_: roboimi.vla.agent_imf.IMFVLAAgent + +action_dim: 16 +obs_dim: 16 +normalization_type: "min_max" +pred_horizon: 16 +obs_horizon: 2 +num_action_steps: 8 +camera_names: ${data.camera_names} +num_cams: ${len:${agent.camera_names}} + +vision_backbone: + num_cameras: ${agent.num_cams} + camera_names: ${agent.camera_names} + vision_backbone: "resnet18" + vision_backbone_mode: "resnet" + freeze_backbone: false + use_separate_rgb_encoder_per_camera: true + output_tokens_per_camera: true + +cond_projector: + output_dim: ${agent.head.n_emb} + +diffusion_steps: 100 +inference_steps: 1 +head_type: "transformer" + +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + cond_dim: ${agent.head.n_emb} + causal_attn: false + time_as_cond: true + obs_as_cond: true + n_cond_layers: 0 + backbone_type: attnres_full + n_head: 1 + n_kv_head: 1 diff --git a/roboimi/vla/conf/agent/resnet_transformer.yaml b/roboimi/vla/conf/agent/resnet_transformer.yaml new file mode 100644 index 0000000..5b129fc --- /dev/null +++ b/roboimi/vla/conf/agent/resnet_transformer.yaml @@ -0,0 +1,62 @@ +# @package agent +defaults: + - /backbone@vision_backbone: resnet_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /head: transformer1d + - _self_ + +_target_: roboimi.vla.agent.VLAAgent + +# ==================== +# 模型维度配置 +# ==================== +action_dim: 16 # 动作维度（机器人关节数） +obs_dim: 16 # 本体感知维度（关节位置） + +# ==================== +# 归一化配置 +# ==================== +normalization_type: "min_max" # "min_max" or "gaussian" + +# ==================== +# 时间步配置 +# ==================== +pred_horizon: 16 # 预测未来多少步动作 +obs_horizon: 2 # 使用多少步历史观测 +num_action_steps: 8 # 每次推理实际执行多少步动作（应 <= pred_horizon - obs_horizon + 1） + +# ==================== +# 相机配置 +# ==================== +camera_names: ${data.camera_names} # 条件相机顺序固定为 r_vis, top, front +num_cams: 3 # 摄像头数量 (r_vis, top, front) + +vision_backbone: + num_cameras: ${agent.num_cams} + camera_names: ${agent.camera_names} + +# ==================== +# 扩散过程配置 +# ==================== +diffusion_steps: 100 # 扩散训练步数（DDPM） +inference_steps: 10 # 推理时的去噪步数（DDIM，��定为 10） + +# ==================== +# Head 类型标识（用于VLAAgent选择调用方式） +# ==================== +head_type: "transformer" # "unet" 或 "transformer" + +# Head 参数覆盖 +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + n_obs_steps: ${agent.obs_horizon} + # Transformer的cond_dim是每步的维度 + # ResNet18 + SpatialSoftmax(32 keypoints) = 64维/相机 + # 计算方式：单相机特征(64) * 相机数(3) + obs_dim(16) = 208 + cond_dim: 208 + causal_attn: false + time_as_cond: true + obs_as_cond: true diff --git a/roboimi/vla/conf/agent/siglip2_imf_attnres.yaml b/roboimi/vla/conf/agent/siglip2_imf_attnres.yaml new file mode 100644 index 0000000..2694621 --- /dev/null +++ b/roboimi/vla/conf/agent/siglip2_imf_attnres.yaml @@ -0,0 +1,44 @@ +# @package agent +defaults: + - /backbone@vision_backbone: siglip2_diffusion + - /modules@state_encoder: identity_state_encoder + - /modules@action_encoder: identity_action_encoder + - /modules@cond_projector: linear_condition_projector + - /head: imf_transformer1d + - _self_ + +_target_: roboimi.vla.agent_imf.IMFVLAAgent + +action_dim: 16 +obs_dim: 16 +normalization_type: "min_max" +pred_horizon: 16 +obs_horizon: 2 +num_action_steps: 8 +camera_names: ${data.camera_names} +num_cams: ${len:${agent.camera_names}} + +vision_backbone: + num_cameras: ${agent.num_cams} + camera_names: ${agent.camera_names} + +cond_projector: + output_dim: ${agent.head.cond_dim} + +diffusion_steps: 100 +inference_steps: 1 +head_type: "transformer" + +head: + input_dim: ${agent.action_dim} + output_dim: ${agent.action_dim} + horizon: ${agent.pred_horizon} + n_obs_steps: ${agent.obs_horizon} + cond_dim: 384 + causal_attn: false + time_as_cond: true + obs_as_cond: true + n_cond_layers: 0 + backbone_type: attnres_full + n_head: 1 + n_kv_head: 1 diff --git a/roboimi/vla/conf/backbone/lewm_vit_diffusion.yaml b/roboimi/vla/conf/backbone/lewm_vit_diffusion.yaml new file mode 100644 index 0000000..9956475 --- /dev/null +++ b/roboimi/vla/conf/backbone/lewm_vit_diffusion.yaml @@ -0,0 +1,16 @@ +_target_: roboimi.vla.models.backbones.lewm_vit_backbone.LEWMViTBackbone + +# LEWM checkpoint path; override this on the target machine. +checkpoint_path: null + +# Input camera contract for roboimi; internal LEWM fusion order stays front/top/r_vis. +num_cameras: 3 +camera_names: [r_vis, top, front] +fused_camera_names: [front, top, r_vis] + +freeze_backbone: true +joint_output_dim: 192 +output_dim: 192 +image_size: 224 +dataset_image_resize_shape: null +eval_image_resize_shape: [256, 256] diff --git a/roboimi/vla/conf/backbone/resnet_diffusion.yaml b/roboimi/vla/conf/backbone/resnet_diffusion.yaml new file mode 100644 index 0000000..23a8d8b --- /dev/null +++ b/roboimi/vla/conf/backbone/resnet_diffusion.yaml @@ -0,0 +1,50 @@ +_target_: roboimi.vla.models.backbones.resnet_diffusion.ResNetDiffusionBackbone + +# ==================== +# 骨干网络选择 +# ==================== +vision_backbone: "resnet18" # torchvision 模型名称: resnet18, resnet34, resnet50 +pretrained_backbone_weights: "IMAGENET1K_V1" # 使用ImageNet预训练权重（torchvision>=0.13） +vision_backbone_mode: "resnet" # resnet | attnres_resnet + +# ==================== +# 冻结设置 +# ==================== +freeze_backbone: true # 冻结ResNet参数，只训练后面的pool和out层（推荐：true） + +# ==================== +# 输入配置 +# ==================== +input_shape: [3, 224, 224] # 输入图像形状 (C, H, W) - ImageNet标准尺寸 +crop_shape: null # 裁剪后的图像形状 (H, W) - 设为null禁用裁剪 +crop_is_random: true # 训练时使用随机裁剪，评估时使用中心裁剪（crop_shape=null时无效） + +# ==================== +# 归一化和特征提取 +# ==================== +use_group_norm: true # 使用 GroupNorm 替代 BatchNorm（更适合小批次训练） +spatial_softmax_num_keypoints: 32 # Spatial Softmax 关键点数量 + +# ==================== +# 编码器模式 +# ==================== +# false: 共享编码器（所有摄像头共享一个 ResNet，参数少但容量受限）推荐！ +# true: 独立编码器（每个摄像头有独立的 ResNet，参数多但容量大） +use_separate_rgb_encoder_per_camera: true +# false: 将所有相机特征拼成一个条件token；true: 每个相机输出一个独立token +output_tokens_per_camera: false +num_cameras: 3 # 摄像头数量 + +# ==================== +# Full-AttnRes vision trunk（当 vision_backbone_mode=attnres_resnet 时生效） +# ==================== +attnres_stem_dim: 64 +attnres_stage_dims: [64, 128, 256, 512] +attnres_stage_depths: [2, 2, 2, 2] +attnres_stage_heads: [4, 4, 8, 8] +attnres_stage_kv_heads: [1, 1, 1, 1] +attnres_stage_window_sizes: [7, 7, 7, 7] +attnres_dropout: 0.0 +attnres_ffn_mult: 2.667 +attnres_eps: 1.0e-06 +attnres_rope_theta: 10000.0 diff --git a/roboimi/vla/conf/backbone/siglip2_diffusion.yaml b/roboimi/vla/conf/backbone/siglip2_diffusion.yaml new file mode 100644 index 0000000..9f0bb62 --- /dev/null +++ b/roboimi/vla/conf/backbone/siglip2_diffusion.yaml @@ -0,0 +1,10 @@ +_target_: roboimi.vla.models.backbones.siglip2_diffusion_backbone.SigLIP2DiffusionBackbone + +model_name: google/siglip2-base-patch16-256 +camera_names: [r_vis, top, front] +num_cameras: 3 +per_view_output_dim: 96 +freeze_backbone: true + +dataset_image_resize_shape: null +eval_image_resize_shape: [256, 256] diff --git a/roboimi/vla/conf/config.yaml b/roboimi/vla/conf/config.yaml new file mode 100644 index 0000000..7f991e0 --- /dev/null +++ b/roboimi/vla/conf/config.yaml @@ -0,0 +1,51 @@ +defaults: + - agent: resnet_transformer + - data: simpe_robot_dataset + - eval: eval + - _self_ + +# ==================== +# 训练配置 +# ==================== +train: + # 基础训练参数 + batch_size: 16 # 批次大小 + lr: 1e-4 # 学习率 + max_steps: 100000 # 最大训练步数 + device: "cuda" # 设备: "cuda" 或 "cpu" + disable_cudnn: false # 遇到当前机器的 cuDNN 兼容性问题时可置 true + + # 数据加载 + num_workers: 12 # DataLoader 工作进程数（调试时设为 0） + val_split: 0.0 # 验证集比例；默认使用全量数据训练 + seed: 42 # 随机种子（用于数据划分） + + # 日志和检查点 + log_freq: 100 # 日志记录频率（步数） + save_freq: 2000 # 保存检查点频率（步数） + use_swanlab: false # 是否启用 SwanLab 标量日志 + swanlab_project: "roboimi-vla" # SwanLab project 名称 + swanlab_run_name: null # 可选的 SwanLab 运行名 + rollout_val_freq_epochs: 50 # 每隔多少个 epoch 执行一次 rollout 验证 + rollout_validate_on_checkpoint: false # 是否在保存 checkpoint 后立即运行 rollout 验证 + rollout_num_episodes: 3 # rollout 验证的回合数 + + # 学习率调度器（带预热） + warmup_steps: 2000 # 预热步数（Transformer建议更长） + scheduler_type: "cosine" # 预热后的调度器: "constant" 或 "cosine" + min_lr: 1e-6 # 最小学习率（用于余弦退火） + + # 优化器 + weight_decay: 1e-5 # 权重衰减（L2 正则化） + grad_clip: 1.0 # 梯度裁剪阈值 + + # 微调配置 + pretrained_ckpt: null # 预训练 checkpoint 路径（用于微调），例如: "checkpoints/vla_model_step_8000.pt" + +# ==================== +# 实验配置 +# ==================== +experiment: + name: "vla_diffusion" # 实验名称 + notes: "" # 实验备注 + tags: [] # 实验标签 diff --git a/roboimi/vla/conf/data/simpe_robot_dataset.yaml b/roboimi/vla/conf/data/simpe_robot_dataset.yaml new file mode 100644 index 0000000..f1af212 --- /dev/null +++ b/roboimi/vla/conf/data/simpe_robot_dataset.yaml @@ -0,0 +1,24 @@ +# @package data +_target_: roboimi.vla.data.simpe_robot_dataset.SimpleRobotDataset + +# ==================== +# 数据集路径 +# ==================== +dataset_dir: "roboimi/demos/dataset/sim_transfer" + +# ==================== +# 时间步参数（从 agent 配置引用） +# ==================== +pred_horizon: ${agent.pred_horizon} # 预测步数 +obs_horizon: ${agent.obs_horizon} # 观测步数 + +# ==================== +# 相机配置 +# ==================== +camera_names: + - r_vis # 机器人视角相机 + - top # 顶部相机 + - front # 前方相机 + +# 单视角预缩放尺寸；为 null 时保留数据集中的原始分辨率 +image_resize_shape: [224, 224] diff --git a/roboimi/vla/conf/eval/eval.yaml b/roboimi/vla/conf/eval/eval.yaml new file mode 100644 index 0000000..06b47e3 --- /dev/null +++ b/roboimi/vla/conf/eval/eval.yaml @@ -0,0 +1,50 @@ +# @package eval +# 评估配置 +ckpt_path: "checkpoints/vla_model_best.pt" # 模型检查点路径 +num_episodes: 3 # 评估回合数 +max_timesteps: 700 # 每回合最大时间步 +device: ${train.device} # 与训练保持一致 +task_name: "sim_transfer" # 环境任务名称 + +# ==================== +# 策略执行参数 +# ==================== +# num_queries 已废弃，现在使用 agent 的 select_action() 自动管理队列 +# 以下参数仅用于兼容旧代码，实际使用 agent.num_action_steps +num_queries: ${agent.num_action_steps} +obs_horizon: ${agent.obs_horizon} + +# ==================== +# 相机配置 +# ==================== +camera_names: ${data.camera_names} + +# ==================== +# 动作平滑 +# ==================== +use_smoothing: false +smooth_method: "ema" +smooth_alpha: 0.3 + +# ==================== +# 调试选项 +# ==================== +headless: false # 是否禁用 MuJoCo / OpenCV GUI 渲染 +verbose_action: true # 是否打印每个时间步的动作信息 + +# ==================== +# Rollout artifact 导出 +# ==================== +artifact_dir: null # 可选输出目录；为空时在启用导出时自动创建目录 +save_artifacts: false # 总开关；实际仍需搭配下面的具体导出项 +save_timing: false # 是否保存 timing.json（包含各阶段耗时统计） +save_trajectory: false # 是否保存 trajectory.npz（原始 EE action + 执行后 EE pose） +save_summary_json: false # 是否保存 JSON-friendly rollout summary +save_trajectory_npz: false # 是否保存每步轨迹/时序/EE pose 为 NPZ +save_trajectory_image: false # 是否保存带红色 EE 轨迹覆盖的静态 PNG +trajectory_image_camera: null # trajectory_image_camera_name 的别名 +trajectory_image_camera_name: null # 导出轨迹图片使用的相机名；为空时默认取 camera_names[0] +record_video: false # 是否从单个相机流录制 rollout mp4 +video_camera: null # video_camera_name 的别名 +video_camera_name: null # 录制视频使用的相机名；为空时默认取 camera_names[0] +video_fps: 30 # 导出 mp4 的目标帧率 diff --git a/roboimi/vla/conf/head/conditional_unet1d.yaml b/roboimi/vla/conf/head/conditional_unet1d.yaml new file mode 100644 index 0000000..b547991 --- /dev/null +++ b/roboimi/vla/conf/head/conditional_unet1d.yaml @@ -0,0 +1,15 @@ +_target_: roboimi.vla.models.heads.conditional_unet1d.ConditionalUnet1D +_partial_: true + +# ==================== +# UNet1D 配置 +# ==================== +kernel_size: 3 # 卷积核大小 +cond_predict_scale: false # FiLM 条件化时是否同时预测 scale（bias + scale 或仅 bias） + +# ==================== +# 网络架构（默认值，可覆盖） +# ==================== +# diffusion_step_embed_dim: 256 # 扩散时间步嵌入维度 +# down_dims: [256, 512, 1024] # 下采样各层通道数 +# n_groups: 8 # GroupNorm 分组数 diff --git a/roboimi/vla/conf/head/gr00t_dit1d.yaml b/roboimi/vla/conf/head/gr00t_dit1d.yaml new file mode 100644 index 0000000..acd0ba7 --- /dev/null +++ b/roboimi/vla/conf/head/gr00t_dit1d.yaml @@ -0,0 +1,22 @@ +_target_: roboimi.vla.models.heads.gr00t_dit1d.Gr00tDiT1D +_partial_: true + +# DiT architecture +n_layer: 6 +n_head: 8 +n_emb: 256 +hidden_dim: 256 +mlp_ratio: 4 +dropout: 0.1 + +# Positional embeddings +add_action_pos_emb: true +add_cond_pos_emb: true + +# Supplied by agent interpolation: +# - input_dim +# - output_dim +# - horizon +# - n_obs_steps +# - cond_dim + diff --git a/roboimi/vla/conf/head/imf_transformer1d.yaml b/roboimi/vla/conf/head/imf_transformer1d.yaml new file mode 100644 index 0000000..92f7054 --- /dev/null +++ b/roboimi/vla/conf/head/imf_transformer1d.yaml @@ -0,0 +1,22 @@ +_target_: roboimi.vla.models.heads.imf_transformer1d.IMFTransformer1D +_partial_: true + +input_dim: ${agent.action_dim} +output_dim: ${agent.action_dim} +horizon: ${agent.pred_horizon} +n_obs_steps: ${agent.obs_horizon} +cond_dim: 208 +n_layer: 12 +n_head: 1 +n_emb: 768 +p_drop_emb: 0.1 +p_drop_attn: 0.1 +causal_attn: false +time_as_cond: true +obs_as_cond: true +n_cond_layers: 0 +backbone_type: attnres_full +n_kv_head: 1 +attn_res_ffn_mult: 2.667 +attn_res_eps: 1.0e-6 +attn_res_rope_theta: 10000.0 diff --git a/roboimi/vla/conf/head/transformer1d.yaml b/roboimi/vla/conf/head/transformer1d.yaml new file mode 100644 index 0000000..4c9cc78 --- /dev/null +++ b/roboimi/vla/conf/head/transformer1d.yaml @@ -0,0 +1,30 @@ +# Transformer-based Diffusion Policy Head +_target_: roboimi.vla.models.heads.transformer1d.Transformer1D +_partial_: true + +# ==================== +# Transformer 架构配置 +# ==================== +n_layer: 4 # Transformer层数（保持当前小模型配置） +n_head: 4 # 注意力头数 +n_emb: 128 # 嵌入维度 +p_drop_emb: 0.05 # Embedding dropout +p_drop_attn: 0.05 # Attention dropout + +# ==================== +# 条件配置 +# ==================== +causal_attn: false # 对齐 external TransformerForDiffusion 的 full-attention / nocausal 变体 +time_as_cond: true # 与 external 实现一致：时间步作为条件 token +obs_as_cond: true # API 对齐；实际是否启用由 cond_dim > 0 决定 +n_cond_layers: 1 # 条件编码器层数（保留当前配置） + +# ==================== +# 注意事项 +# ==================== +# 以下参数将在agent配置中通过interpolation提供： +# - input_dim: ${agent.action_dim} +# - output_dim: ${agent.action_dim} +# - horizon: ${agent.pred_horizon} +# - n_obs_steps: ${agent.obs_horizon} +# - cond_dim: 通过agent中的global_cond_dim计算 diff --git a/roboimi/vla/conf/modules/identity_action_encoder.yaml b/roboimi/vla/conf/modules/identity_action_encoder.yaml new file mode 100644 index 0000000..4f18b51 --- /dev/null +++ b/roboimi/vla/conf/modules/identity_action_encoder.yaml @@ -0,0 +1 @@ +_target_: roboimi.vla.modules.encoders.IdentityActionEncoder diff --git a/roboimi/vla/conf/modules/identity_state_encoder.yaml b/roboimi/vla/conf/modules/identity_state_encoder.yaml new file mode 100644 index 0000000..fba00d5 --- /dev/null +++ b/roboimi/vla/conf/modules/identity_state_encoder.yaml @@ -0,0 +1 @@ +_target_: roboimi.vla.modules.encoders.IdentityStateEncoder diff --git a/roboimi/vla/conf/modules/linear_condition_projector.yaml b/roboimi/vla/conf/modules/linear_condition_projector.yaml new file mode 100644 index 0000000..c5565b1 --- /dev/null +++ b/roboimi/vla/conf/modules/linear_condition_projector.yaml @@ -0,0 +1,5 @@ +_target_: roboimi.vla.modules.projectors.LinearConditionProjector +_partial_: true + +output_dim: 384 +bias: true diff --git a/roboimi/vla/core/__init__.py b/roboimi/vla/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboimi/vla/core/interfaces.py b/roboimi/vla/core/interfaces.py new file mode 100644 index 0000000..ea02094 --- /dev/null +++ b/roboimi/vla/core/interfaces.py @@ -0,0 +1,46 @@ +import abc +import torch +import torch.nn as nn +from typing import Dict, Any, Optional + +class VLABackbone(nn.Module, abc.ABC): + """ + Contract for Vision/Language Backbones. + Must return a feature tensor of shape (B, Seq, Embed_Dim). + """ + @abc.abstractmethod + def forward(self, obs: Dict[str, torch.Tensor]) -> torch.Tensor: + """ + Args: + obs: Dictionary containing 'image' and optionally 'text'. + Returns: + features: (B, S, D) embedding. + """ + pass + + +class VLAProjector(nn.Module, abc.ABC): + """ + Contract for the adaptation layer (Projector). + Connects Backbone features to the Policy Head. + """ + @abc.abstractmethod + def forward(self, x: torch.Tensor) -> torch.Tensor: + pass + + +class VLAHead(nn.Module, abc.ABC): + """ + Contract for Action Generation Heads (Policies). + Handles both training (loss calculation) and inference (action generation). + """ + @abc.abstractmethod + def forward(self, embeddings: torch.Tensor, actions: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + """ + Args: + embeddings: (B, S, Hidden) from Projector. + actions: (B, Pred_Horizon, Action_Dim) - Ground truth for training. + Returns: + Dict containing 'loss' (if actions provided) or 'pred_actions'. + """ + pass \ No newline at end of file diff --git a/roboimi/vla/data/__init__.py b/roboimi/vla/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboimi/vla/data/simpe_robot_dataset.py b/roboimi/vla/data/simpe_robot_dataset.py new file mode 100644 index 0000000..c1c5a93 --- /dev/null +++ b/roboimi/vla/data/simpe_robot_dataset.py @@ -0,0 +1,249 @@ +import torch +import h5py +from torch.utils.data import Dataset +from typing import List, Dict, Union, Optional, Sequence +from pathlib import Path +from collections import OrderedDict + + +class SimpleRobotDataset(Dataset): + """ + HDF5 懒加载数据集 - LeRobotDataset 格式 + + 返回格式: + - observation.state: (obs_horizon, state_dim) + - observation.{cam_name}: (obs_horizon, C, H, W) + - action: (pred_horizon, action_dim) + """ + + def __init__( + self, + dataset_dir: Union[str, Path], + obs_horizon: int = 2, + pred_horizon: int = 8, + camera_names: List[str] = None, + image_resize_shape: Optional[Sequence[int]] = (224, 224), + max_open_files: int = 64, + ): + """ + Args: + dataset_dir: HDF5 文件目录路径 + obs_horizon: 观察过去多少帧 + pred_horizon: 预测未来多少帧动作 + camera_names: 相机名称列表，如 ["r_vis", "top", "front"] + image_resize_shape: 图像缩放尺寸 (W, H)；为 None 时保留原始分辨率 + max_open_files: 每个 worker 最多缓存的 HDF5 文件句柄数 + + HDF5 文件格式: + - action: [T, action_dim] + - observations/qpos: [T, obs_dim] + - observations/images/{cam_name}: [T, H, W, C] + """ + self.obs_horizon = obs_horizon + self.pred_horizon = pred_horizon + self.camera_names = camera_names or [] + self.image_resize_shape = ( + tuple(int(v) for v in image_resize_shape) + if image_resize_shape is not None else None + ) + self.max_open_files = max(1, int(max_open_files)) + self._file_cache: "OrderedDict[str, h5py.File]" = OrderedDict() + + self.dataset_dir = Path(dataset_dir) + if not self.dataset_dir.exists(): + raise FileNotFoundError(f"数据集目录不存在: {dataset_dir}") + + # 查找 HDF5 文件 + self.hdf5_files = sorted(self.dataset_dir.glob("*.hdf5")) + if not self.hdf5_files: + self.hdf5_files = sorted(self.dataset_dir.glob("episode_*.hdf5")) + if not self.hdf5_files: + raise FileNotFoundError(f"在 {dataset_dir} 中未找到 HDF5 文件") + + # 构建 episode 索引（只存储元数据，不加载数据） + self.episodes = {} + self.frame_meta = [] # 存储 (ep_idx, frame_idx, hdf5_path) + for ep_idx, hdf5_path in enumerate(self.hdf5_files): + with h5py.File(hdf5_path, 'r') as f: + T = f['action'].shape[0] + start_idx = len(self.frame_meta) + for t in range(T): + self.frame_meta.append({ + "ep_idx": ep_idx, + "frame_idx": t, + "hdf5_path": hdf5_path, + }) + self.episodes[ep_idx] = list(range(start_idx, len(self.frame_meta))) + + print(f"懒加载模式: {len(self.hdf5_files)} 个 episodes, 共 {len(self.frame_meta)} 帧") + + def __len__(self): + return len(self.frame_meta) + + def _close_all_files(self) -> None: + """关闭当前 worker 内缓存的所有 HDF5 文件句柄。""" + for f in self._file_cache.values(): + try: + f.close() + except Exception: + pass + self._file_cache.clear() + + def _get_h5_file(self, hdf5_path: Union[str, Path]) -> h5py.File: + """ + 获取 HDF5 文件句柄（worker 内 LRU 缓存）。 + 注意：缓存的是文件句柄，不是帧数据本身。 + """ + key = str(hdf5_path) + if key in self._file_cache: + self._file_cache.move_to_end(key) + return self._file_cache[key] + + # 超过上限时淘汰最久未使用的句柄 + if len(self._file_cache) >= self.max_open_files: + _, old_file = self._file_cache.popitem(last=False) + try: + old_file.close() + except Exception: + pass + + f = h5py.File(key, 'r') + self._file_cache[key] = f + return f + + def _load_frame(self, idx: int, *, load_images: bool = True) -> Dict: + """从 HDF5 文件懒加载单帧数据""" + meta = self.frame_meta[idx] + f = self._get_h5_file(meta["hdf5_path"]) + frame = { + "episode_index": meta["ep_idx"], + "frame_index": meta["frame_idx"], + "task": f.get('task', [b"unknown"])[0].decode() if 'task' in f else "unknown", + "observation.state": torch.from_numpy(f['observations/qpos'][meta["frame_idx"]]).float(), + "action": torch.from_numpy(f['action'][meta["frame_idx"]]).float(), + } + + # 加载图像数据: observations/images/{cam_name} -> observation.{cam_name} + if load_images: + for cam_name in self.camera_names: + h5_path = f'observations/images/{cam_name}' + if h5_path in f: + img = f[h5_path][meta["frame_idx"]] + if self.image_resize_shape is not None: + import cv2 + img = cv2.resize(img, self.image_resize_shape, interpolation=cv2.INTER_LINEAR) + # 转换为float并归一化到 [0, 1] + img = torch.from_numpy(img).float() / 255.0 + frame[f"observation.{cam_name}"] = img.permute(2, 0, 1) # HWC -> CHW + + return frame + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + frame = self._load_frame(idx, load_images=False) + ep_idx = frame["episode_index"] + + # 获取当前 episode 的帧索引范围 + ep_indices = self.episodes[ep_idx] + ep_start = ep_indices[0] + ep_end = ep_indices[-1] + + # ============================================ + # 1. 加载观察（过去 obs_horizon 帧） + # ============================================ + observations = { + "state": [], # 状态数据 + } + # 为每个摄像头初始化独立列表 + for cam_name in self.camera_names: + observations[f"observation.{cam_name}"] = [] + + observation_is_pad = [] + + for delta in range(-self.obs_horizon + 1, 1): # [-1, 0] for obs_horizon=2 + target_idx = idx + delta + + # 边界检查 + if ep_start <= target_idx <= ep_end: + target_frame = self._load_frame(target_idx) + is_pad = False + else: + # 超出边界，用边界帧填充 + if target_idx < ep_start: + target_frame = self._load_frame(ep_start) + else: + target_frame = self._load_frame(ep_end) + is_pad = True + + # 收集状态 + observations["state"].append(target_frame["observation.state"]) + + # 收集每个摄像头的图像 + for cam_name in self.camera_names: + observations[f"observation.{cam_name}"].append(target_frame[f"observation.{cam_name}"]) + + observation_is_pad.append(is_pad) + + # ============================================ + # 2. 加载动作（未来 pred_horizon 帧） + # ============================================ + actions = [] + action_is_pad = [] + + for delta in range(self.pred_horizon): + target_idx = idx + delta + + if target_idx <= ep_end: + actions.append(self._load_frame(target_idx, load_images=False)["action"]) + action_is_pad.append(False) + else: + actions.append(self._load_frame(ep_end, load_images=False)["action"]) + action_is_pad.append(True) + + # ============================================ + # 3. 组装返回数据（LeRobotDataset 格式） + # ============================================ + result = { + # 状态观察: (obs_horizon, state_dim) + "observation.state": torch.stack(observations["state"]), + "observation_is_pad": torch.tensor(observation_is_pad, dtype=torch.bool), + + # 动作: (pred_horizon, action_dim) + "action": torch.stack(actions), + "action_is_pad": torch.tensor(action_is_pad, dtype=torch.bool), + + # 任务 + "task": frame["task"], + } + + # 图像：每个摄像头独立的 key + # 形状: (obs_horizon, C, H, W) + for cam_name in self.camera_names: + result[f"observation.{cam_name}"] = torch.stack(observations[f"observation.{cam_name}"]) + + return result + + @property + def camera_keys(self) -> list[str]: + """获取所有相机键名 (LeRobotDataset 格式)""" + return [f"observation.{cam_name}" for cam_name in self.camera_names] + + @property + def camera_info(self) -> dict: + """获取相机信息""" + if not self.camera_names: + return {} + + # 从第一个样本获取形状 + sample = self[0] + info = {} + for cam_name in self.camera_names: + key = f"observation.{cam_name}" + if key in sample: + info[key] = { + "shape": sample[key].shape, + "dtype": str(sample[key].dtype), + } + return info + + def __del__(self): + self._close_all_files() diff --git a/roboimi/vla/eval_utils.py b/roboimi/vla/eval_utils.py new file mode 100644 index 0000000..73cb05d --- /dev/null +++ b/roboimi/vla/eval_utils.py @@ -0,0 +1,3 @@ +def execute_policy_action(env, action): + """Execute policy outputs using EE-action semantics.""" + env.step(action) diff --git a/roboimi/vla/models/__init__.py b/roboimi/vla/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboimi/vla/models/backbones/__init__.py b/roboimi/vla/models/backbones/__init__.py new file mode 100644 index 0000000..c5544b5 --- /dev/null +++ b/roboimi/vla/models/backbones/__init__.py @@ -0,0 +1,15 @@ +# Backbone models +__all__ = ["LEWMViTBackbone", "ResNetBackbone", "ResNetDiffusionBackbone", "SigLIP2DiffusionBackbone"] + + +def __getattr__(name): + if name == "LEWMViTBackbone": + from .lewm_vit_backbone import LEWMViTBackbone + return LEWMViTBackbone + if name == "SigLIP2DiffusionBackbone": + from .siglip2_diffusion_backbone import SigLIP2DiffusionBackbone + return SigLIP2DiffusionBackbone + if name in {"ResNetBackbone", "ResNetDiffusionBackbone"}: + from .resnet_diffusion import ResNetDiffusionBackbone + return ResNetDiffusionBackbone + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/roboimi/vla/models/backbones/attnres_resnet2d.py b/roboimi/vla/models/backbones/attnres_resnet2d.py new file mode 100644 index 0000000..1ef144f --- /dev/null +++ b/roboimi/vla/models/backbones/attnres_resnet2d.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +from typing import Iterable, Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from roboimi.vla.models.heads.attnres_transformer_components import AttnResTransformerBackbone + + +def _make_norm2d(num_channels: int, use_group_norm: bool) -> nn.Module: + if use_group_norm: + num_groups = max(1, num_channels // 16) + return nn.GroupNorm(num_groups=num_groups, num_channels=num_channels) + return nn.BatchNorm2d(num_channels) + + +class _ConvNormAct(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + *, + kernel_size: int, + stride: int, + padding: int, + use_group_norm: bool, + ) -> None: + super().__init__() + self.block = nn.Sequential( + nn.Conv2d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=False, + ), + _make_norm2d(out_channels, use_group_norm), + nn.SiLU(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.block(x) + + +class AttnResImageBlock2D(nn.Module): + def __init__( + self, + dim: int, + *, + window_size: int, + n_heads: int, + n_kv_heads: int, + dropout: float, + ffn_mult: float, + eps: float, + rope_theta: float, + ) -> None: + super().__init__() + self.window_size = int(window_size) + self.block = AttnResTransformerBackbone( + d_model=dim, + n_blocks=1, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + max_seq_len=self.window_size * self.window_size, + dropout=dropout, + ffn_mult=ffn_mult, + eps=eps, + rope_theta=rope_theta, + causal_attn=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + bsz, channels, height, width = x.shape + ws = self.window_size + pad_h = (ws - height % ws) % ws + pad_w = (ws - width % ws) % ws + if pad_h or pad_w: + x = F.pad(x, (0, pad_w, 0, pad_h)) + padded_height, padded_width = x.shape[-2:] + num_h = padded_height // ws + num_w = padded_width // ws + + windows = ( + x.permute(0, 2, 3, 1) + .contiguous() + .view(bsz, num_h, ws, num_w, ws, channels) + .permute(0, 1, 3, 2, 4, 5) + .contiguous() + .view(bsz * num_h * num_w, ws * ws, channels) + ) + windows = self.block(windows) + x = ( + windows.view(bsz, num_h, num_w, ws, ws, channels) + .permute(0, 1, 3, 2, 4, 5) + .contiguous() + .view(bsz, padded_height, padded_width, channels) + .permute(0, 3, 1, 2) + .contiguous() + ) + return x[:, :, :height, :width] + + +class _AttnResStage2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + *, + depth: int, + downsample_stride: int, + use_group_norm: bool, + window_size: int, + n_heads: int, + n_kv_heads: int, + dropout: float, + ffn_mult: float, + eps: float, + rope_theta: float, + ) -> None: + super().__init__() + self.downsample = None + if in_channels != out_channels or downsample_stride != 1: + kernel_size = 1 if downsample_stride == 1 else 3 + padding = 0 if downsample_stride == 1 else 1 + self.downsample = _ConvNormAct( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=downsample_stride, + padding=padding, + use_group_norm=use_group_norm, + ) + + self.blocks = nn.ModuleList( + [ + AttnResImageBlock2D( + out_channels, + window_size=window_size, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + dropout=dropout, + ffn_mult=ffn_mult, + eps=eps, + rope_theta=rope_theta, + ) + for _ in range(depth) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.downsample is not None: + x = self.downsample(x) + for block in self.blocks: + x = block(x) + return x + + +class AttnResResNetLikeBackbone2D(nn.Module): + def __init__( + self, + *, + input_channels: int = 3, + stem_dim: int = 64, + stage_dims: Sequence[int] = (64, 128, 256, 512), + stage_depths: Sequence[int] = (2, 2, 2, 2), + stage_heads: Sequence[int] = (4, 4, 8, 8), + stage_kv_heads: Sequence[int] = (1, 1, 1, 1), + stage_window_sizes: Sequence[int] = (7, 7, 7, 7), + use_group_norm: bool = True, + dropout: float = 0.0, + ffn_mult: float = 2.667, + eps: float = 1e-6, + rope_theta: float = 10000.0, + ) -> None: + super().__init__() + + lengths = { + len(stage_dims), + len(stage_depths), + len(stage_heads), + len(stage_kv_heads), + len(stage_window_sizes), + } + if len(lengths) != 1: + raise ValueError('stage_dims/depths/heads/kv_heads/window_sizes 长度必须一致') + + self.stem = nn.Sequential( + nn.Conv2d(input_channels, stem_dim, kernel_size=7, stride=2, padding=3, bias=False), + _make_norm2d(stem_dim, use_group_norm), + nn.SiLU(), + nn.MaxPool2d(kernel_size=3, stride=2, padding=1), + ) + + in_channels = stem_dim + stages = [] + for stage_idx, (out_channels, depth, n_heads, n_kv_heads, window_size) in enumerate( + zip(stage_dims, stage_depths, stage_heads, stage_kv_heads, stage_window_sizes) + ): + stages.append( + _AttnResStage2D( + in_channels=in_channels, + out_channels=out_channels, + depth=int(depth), + downsample_stride=1 if stage_idx == 0 else 2, + use_group_norm=use_group_norm, + window_size=int(window_size), + n_heads=int(n_heads), + n_kv_heads=int(n_kv_heads), + dropout=float(dropout), + ffn_mult=float(ffn_mult), + eps=float(eps), + rope_theta=float(rope_theta), + ) + ) + in_channels = out_channels + + self.stages = nn.ModuleList(stages) + self.output_channels = in_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.stem(x) + for stage in self.stages: + x = stage(x) + return x diff --git a/roboimi/vla/models/backbones/lewm_vit_backbone.py b/roboimi/vla/models/backbones/lewm_vit_backbone.py new file mode 100644 index 0000000..3ea4683 --- /dev/null +++ b/roboimi/vla/models/backbones/lewm_vit_backbone.py @@ -0,0 +1,230 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Mapping, Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from roboimi.vla.core.interfaces import VLABackbone + + +class _LEWMProjector(nn.Module): + """LEWM projector MLP: 192 -> 2048 -> 192 with BatchNorm1d + GELU.""" + + def __init__(self, input_dim: int = 192, hidden_dim: int = 2048, output_dim: int = 192) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.BatchNorm1d(hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, output_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class LEWMViTBackbone(VLABackbone): + """Frozen LEWM joint-multiview ViT backbone. + + The backbone fuses the three camera views into a single LEWM-style image, + runs a ViT-tiny encoder plus the LEWM projector, and returns one joint + 192-d embedding per timestep. + """ + + def __init__( + self, + checkpoint_path: str | Path | None = None, + *, + checkpoint: Mapping[str, Any] | None = None, + camera_names: Sequence[str] = ("r_vis", "top", "front"), + fused_camera_names: Sequence[str] = ("front", "top", "r_vis"), + num_cameras: int | None = None, + dataset_image_resize_shape: Sequence[int] | None = None, + eval_image_resize_shape: Sequence[int] | None = (256, 256), + freeze_backbone: bool = True, + joint_output_dim: int = 192, + image_size: int = 224, + output_dim: int = 192, + ) -> None: + super().__init__() + + self.camera_names = tuple(camera_names) + self.fused_camera_names = tuple(fused_camera_names) + self.num_cameras = int(num_cameras) if num_cameras is not None else len(self.camera_names) + self.freeze_backbone = bool(freeze_backbone) + self.joint_output_dim = int(joint_output_dim) + self.image_size = int(image_size) + self._output_dim = int(output_dim) + self.dataset_image_resize_shape = ( + tuple(int(v) for v in dataset_image_resize_shape) + if dataset_image_resize_shape is not None else None + ) + self.eval_image_resize_shape = ( + tuple(int(v) for v in eval_image_resize_shape) + if eval_image_resize_shape is not None else None + ) + if self.num_cameras != len(self.camera_names): + raise ValueError( + f"num_cameras({self.num_cameras}) must match len(camera_names)({len(self.camera_names)})" + ) + if set(self.fused_camera_names) != set(self.camera_names): + raise ValueError( + "fused_camera_names must contain the same cameras as camera_names. " + f"got camera_names={list(self.camera_names)}, fused_camera_names={list(self.fused_camera_names)}" + ) + + self.encoder = self._build_encoder(self.image_size) + self.projector = _LEWMProjector( + input_dim=self.encoder.config.hidden_size, + hidden_dim=2048, + output_dim=self.joint_output_dim, + ) + + self.register_buffer( + "mean", + torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(1, 3, 1, 1), + ) + self.register_buffer( + "std", + torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(1, 3, 1, 1), + ) + + if checkpoint_path is not None and checkpoint is not None: + raise ValueError("checkpoint_path and checkpoint cannot both be provided") + if checkpoint_path is not None: + self.load_lewm_checkpoint(checkpoint_path) + elif checkpoint is not None: + self.load_lewm_checkpoint(checkpoint) + + if self.freeze_backbone: + self._freeze_encoder_and_projector() + + @staticmethod + def _build_encoder_config(image_size: int): + from transformers import ViTConfig + + return ViTConfig( + image_size=image_size, + patch_size=14, + num_channels=3, + hidden_size=192, + intermediate_size=768, + num_hidden_layers=12, + num_attention_heads=3, + qkv_bias=True, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + ) + + @classmethod + def _build_encoder(cls, image_size: int) -> nn.Module: + from transformers import ViTModel + + return ViTModel(cls._build_encoder_config(image_size), add_pooling_layer=False) + + @staticmethod + def _unwrap_state_dict(payload: Mapping[str, Any]) -> Mapping[str, torch.Tensor]: + state_dict = payload.get("state_dict", payload) + if not isinstance(state_dict, Mapping): + raise TypeError("checkpoint payload must contain a mapping state_dict") + return state_dict + + @staticmethod + def _extract_prefixed_state_dict( + state_dict: Mapping[str, torch.Tensor], + prefix: str, + ) -> Dict[str, torch.Tensor]: + extracted = { + key[len(prefix) :]: value + for key, value in state_dict.items() + if key.startswith(prefix) + } + if not extracted: + raise KeyError(f"checkpoint missing parameters with prefix {prefix!r}") + return extracted + + def load_lewm_checkpoint(self, checkpoint_or_path: str | Path | Mapping[str, Any]) -> None: + if isinstance(checkpoint_or_path, (str, Path)): + payload = torch.load(Path(checkpoint_or_path), map_location="cpu", weights_only=False) + else: + payload = checkpoint_or_path + + state_dict = self._unwrap_state_dict(payload) + encoder_state_dict = self._extract_prefixed_state_dict(state_dict, "model.encoder.") + projector_state_dict = self._extract_prefixed_state_dict(state_dict, "model.projector.") + + self.encoder.load_state_dict(encoder_state_dict, strict=True) + self.projector.load_state_dict(projector_state_dict, strict=True) + + def _freeze_encoder_and_projector(self) -> None: + for module in (self.encoder, self.projector): + module.eval() + for parameter in module.parameters(): + parameter.requires_grad = False + + def train(self, mode: bool = True) -> "LEWMViTBackbone": + super().train(mode) + if self.freeze_backbone: + self._freeze_encoder_and_projector() + return self + + def _ordered_images(self, images: Dict[str, torch.Tensor]) -> list[torch.Tensor]: + missing = [camera_name for camera_name in self.camera_names if camera_name not in images] + if missing: + raise ValueError( + f"image input missing required cameras. missing={missing}, expected={list(self.camera_names)}" + ) + + ordered = [images[camera_name] for camera_name in self.camera_names] + reference_shape = ordered[0].shape + if len(reference_shape) != 5: + raise ValueError(f"expected image tensors shaped (B, T, C, H, W), got {reference_shape}") + + for camera_name, image in zip(self.camera_names[1:], ordered[1:]): + if image.shape != reference_shape: + raise ValueError( + f"camera {camera_name!r} shape {tuple(image.shape)} does not match {tuple(reference_shape)}" + ) + + return ordered + + def _prepare_pixels(self, images: Dict[str, torch.Tensor]) -> tuple[torch.Tensor, int, int]: + self._ordered_images(images) + fused = torch.cat([images[camera_name] for camera_name in self.fused_camera_names], dim=-2) + bsz, steps = fused.shape[:2] + fused = fused.reshape(bsz * steps, *fused.shape[2:]).contiguous().float() + + fused = fused.clamp(0.0, 1.0) + fused = (fused - self.mean) / self.std + + height, width = fused.shape[-2:] + short_side = min(height, width) + if short_side <= 0: + raise ValueError(f"invalid fused image shape: {tuple(fused.shape)}") + scale = self.image_size / float(short_side) + resized_height = int(round(height * scale)) + resized_width = int(round(width * scale)) + if (resized_height, resized_width) != (height, width): + fused = F.interpolate( + fused, + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + antialias=True, + ) + return fused, bsz, steps + + def forward(self, images: Dict[str, torch.Tensor]) -> torch.Tensor: + pixels, bsz, steps = self._prepare_pixels(images) + with torch.set_grad_enabled(torch.is_grad_enabled() and not self.freeze_backbone): + output = self.encoder(pixel_values=pixels, interpolate_pos_encoding=True) + cls = output.last_hidden_state[:, 0] + embedding = self.projector(cls) + return embedding.view(bsz, steps, self.joint_output_dim) + + @property + def output_dim(self) -> int: + return self._output_dim diff --git a/roboimi/vla/models/backbones/resnet_diffusion.py b/roboimi/vla/models/backbones/resnet_diffusion.py new file mode 100644 index 0000000..00e2d0b --- /dev/null +++ b/roboimi/vla/models/backbones/resnet_diffusion.py @@ -0,0 +1,466 @@ +from roboimi.vla.core.interfaces import VLABackbone +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision +import numpy as np +from typing import Callable, Optional, Tuple, Union + +from .attnres_resnet2d import AttnResResNetLikeBackbone2D + +def _replace_submodules( + root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module] +) -> nn.Module: + """ + Args: + root_module: 需要替换子模块的根模块 + predicate: 接受一个模块作为参数，如果该模块需要被替换则返回 True。 + func: 接受一个模块作为参数，并返回一个新的模块来替换它。 + Returns: + 子模块已被替换的根模块。 + """ + if predicate(root_module): + return func(root_module) + + replace_list = [k.split(".") for k, m in root_module.named_modules(remove_duplicate=True) if predicate(m)] + for *parents, k in replace_list: + parent_module = root_module + if len(parents) > 0: + parent_module = root_module.get_submodule(".".join(parents)) + if isinstance(parent_module, nn.Sequential): + src_module = parent_module[int(k)] + else: + src_module = getattr(parent_module, k) + tgt_module = func(src_module) + if isinstance(parent_module, nn.Sequential): + parent_module[int(k)] = tgt_module + else: + setattr(parent_module, k, tgt_module) + # 验证所有 BN 是否已被替换 + assert not any(predicate(m) for _, m in root_module.named_modules(remove_duplicate=True)) + return root_module + +class SpatialSoftmax(nn.Module): + """ + Finn 等人在 "Deep Spatial Autoencoders for Visuomotor Learning" 中描述的空间软 Argmax 操作 + (https://huggingface.co/papers/1509.06113)。这是 robomimic 实现的一个最小移植版本。 + """ + + def __init__(self, input_shape, num_kp=None): + """ + Args: + input_shape (list): (C, H, W) 输入特征图形状。 + num_kp (int): 输出中的关键点数量。如果为 None，输出将具有与输入相同的通道数。 + """ + super().__init__() + + assert len(input_shape) == 3 + self._in_c, self._in_h, self._in_w = input_shape + + if num_kp is not None: + self.nets = torch.nn.Conv2d(self._in_c, num_kp, kernel_size=1) + self._out_c = num_kp + else: + self.nets = None + self._out_c = self._in_c + + # 我们可以直接使用 torch.linspace，但这似乎与 numpy 的行为略有不同 + # 并且会导致预训练模型的 pc_success 略有下降。 + pos_x, pos_y = np.meshgrid(np.linspace(-1.0, 1.0, self._in_w), np.linspace(-1.0, 1.0, self._in_h)) + pos_x = torch.from_numpy(pos_x.reshape(self._in_h * self._in_w, 1)).float() + pos_y = torch.from_numpy(pos_y.reshape(self._in_h * self._in_w, 1)).float() + # 注册为 buffer，以便将其移动到正确的设备。 + self.register_buffer("pos_grid", torch.cat([pos_x, pos_y], dim=1)) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """ + Args: + features: (B, C, H, W) 输入特征图。 + Returns: + (B, K, 2) 关键点的图像空间坐标。 + """ + if self.nets is not None: + features = self.nets(features) + + # [B, K, H, W] -> [B * K, H * W]，其中 K 是关键点数量 + features = features.reshape(-1, self._in_h * self._in_w) + # 2d softmax 归一化 + attention = F.softmax(features, dim=-1) + # [B * K, H * W] x [H * W, 2] -> [B * K, 2] 用于 x 和 y 维度的空间坐标均值 + expected_xy = attention @ self.pos_grid + # 重塑为 [B, K, 2] + feature_keypoints = expected_xy.view(-1, self._out_c, 2) + + return feature_keypoints + +class _SingleRgbEncoder(nn.Module): + """单个摄像头的 RGB 编码器，支持独立或共享使用""" + def __init__( + self, + vision_backbone: str, + pretrained_backbone_weights: str | None, + input_shape: Tuple[int, int, int], + crop_shape: Optional[Tuple[int, int]], + crop_is_random: bool, + use_group_norm: bool, + spatial_softmax_num_keypoints: int, + freeze_backbone: bool = True, # 新增：是否冻结backbone + vision_backbone_mode: str = "resnet", + attnres_stem_dim: int = 64, + attnres_stage_dims: Optional[Tuple[int, ...]] = None, + attnres_stage_depths: Optional[Tuple[int, ...]] = None, + attnres_stage_heads: Optional[Tuple[int, ...]] = None, + attnres_stage_kv_heads: Optional[Tuple[int, ...]] = None, + attnres_stage_window_sizes: Optional[Tuple[int, ...]] = None, + attnres_dropout: float = 0.0, + attnres_ffn_mult: float = 2.667, + attnres_eps: float = 1e-6, + attnres_rope_theta: float = 10000.0, + ): + super().__init__() + + # 设置可选的预处理 + if crop_shape is not None: + self.do_crop = True + # 评估时始终使用中心裁剪 + self.center_crop = torchvision.transforms.CenterCrop(crop_shape) + if crop_is_random: + self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape) + else: + self.maybe_random_crop = self.center_crop + else: + self.do_crop = False + crop_shape = input_shape[1:] + + if vision_backbone_mode == "resnet": + # 设置骨干网络 + backbone_model = getattr(torchvision.models, vision_backbone)( + weights=pretrained_backbone_weights + ) + + # 移除 AvgPool 和 FC (假设 layer4 是 children()[-3]) + self.backbone = nn.Sequential(*(list(backbone_model.children())[:-2])) + + if use_group_norm: + self.backbone = _replace_submodules( + root_module=self.backbone, + predicate=lambda x: isinstance(x, nn.BatchNorm2d), + func=lambda x: nn.GroupNorm( + num_groups=max(1, x.num_features // 16), + num_channels=x.num_features, + ), + ) + elif vision_backbone_mode == "attnres_resnet": + self.backbone = AttnResResNetLikeBackbone2D( + input_channels=input_shape[0], + stem_dim=attnres_stem_dim, + stage_dims=tuple(attnres_stage_dims or (64, 128, 256, 512)), + stage_depths=tuple(attnres_stage_depths or (2, 2, 2, 2)), + stage_heads=tuple(attnres_stage_heads or (4, 4, 8, 8)), + stage_kv_heads=tuple(attnres_stage_kv_heads or (1, 1, 1, 1)), + stage_window_sizes=tuple(attnres_stage_window_sizes or (7, 7, 7, 7)), + use_group_norm=use_group_norm, + dropout=attnres_dropout, + ffn_mult=attnres_ffn_mult, + eps=attnres_eps, + rope_theta=attnres_rope_theta, + ) + else: + raise ValueError(f"不支持的 vision_backbone_mode: {vision_backbone_mode}") + + # 冻结backbone参数（可选） + if freeze_backbone: + for param in self.backbone.parameters(): + param.requires_grad = False + + # 设置池化和最终层 + # 使用试运行来获取特征图形状 + dummy_shape = (1, input_shape[0], *crop_shape) + with torch.no_grad(): + dummy_out = self.backbone(torch.zeros(dummy_shape)) + feature_map_shape = dummy_out.shape[1:] # (C, H, W) + + self.pool = SpatialSoftmax(feature_map_shape, num_kp=spatial_softmax_num_keypoints) + self.feature_dim = spatial_softmax_num_keypoints * 2 + self.out = nn.Linear(spatial_softmax_num_keypoints * 2, self.feature_dim) + self.relu = nn.ReLU() + + # 注册ImageNet标准化参数为buffer（会自动移到GPU） + self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) + self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) + + def forward_single_image(self, x: torch.Tensor) -> torch.Tensor: + if self.do_crop: + x = self.maybe_random_crop(x) if self.training else self.center_crop(x) + + # ImageNet标准化（预训练权重期望的输入分布） + x = (x - self.mean) / self.std + + x = self.relu(self.out(torch.flatten(self.pool(self.backbone(x)), start_dim=1))) + return x + + +class ResNetDiffusionBackbone(VLABackbone): + def __init__( + self, + vision_backbone: str = "resnet18", + pretrained_backbone_weights: str | None = None, + input_shape: Tuple[int, int, int] = (3, 84, 84), # (C, H, W) + crop_shape: Optional[Tuple[int, int]] = None, + crop_is_random: bool = True, + use_group_norm: bool = True, + spatial_softmax_num_keypoints: int = 32, + use_separate_rgb_encoder_per_camera: bool = False, # 新增：是否为每个摄像头使用独立编码器 + output_tokens_per_camera: bool = False, # 是否按相机返回多token，而不是拼成一个token + num_cameras: int = 1, # 新增：摄像头数量（仅在独立编码器模式下使用） + camera_names: Optional[Tuple[str, ...]] = None, # 显式相机顺序 + freeze_backbone: bool = True, # 新增：是否冻结ResNet backbone（推荐True） + vision_backbone_mode: str = "resnet", + attnres_stem_dim: int = 64, + attnres_stage_dims: Optional[Tuple[int, ...]] = None, + attnres_stage_depths: Optional[Tuple[int, ...]] = None, + attnres_stage_heads: Optional[Tuple[int, ...]] = None, + attnres_stage_kv_heads: Optional[Tuple[int, ...]] = None, + attnres_stage_window_sizes: Optional[Tuple[int, ...]] = None, + attnres_dropout: float = 0.0, + attnres_ffn_mult: float = 2.667, + attnres_eps: float = 1e-6, + attnres_rope_theta: float = 10000.0, + ): + super().__init__() + + self.use_separate_rgb_encoder_per_camera = use_separate_rgb_encoder_per_camera + self.output_tokens_per_camera = bool(output_tokens_per_camera) + self.num_cameras = num_cameras + self.tokens_per_step = self.num_cameras if self.output_tokens_per_camera else 1 + self.camera_names = tuple(camera_names) if camera_names is not None else None + if self.camera_names is not None and len(self.camera_names) != self.num_cameras: + raise ValueError( + f"camera_names 长度({len(self.camera_names)})与 num_cameras({self.num_cameras})不一致" + ) + + if use_separate_rgb_encoder_per_camera: + # 独立编码器模式：为每个摄像头创建独立的编码器 + encoders = [ + _SingleRgbEncoder( + vision_backbone=vision_backbone, + pretrained_backbone_weights=pretrained_backbone_weights, + input_shape=input_shape, + crop_shape=crop_shape, + crop_is_random=crop_is_random, + use_group_norm=use_group_norm, + spatial_softmax_num_keypoints=spatial_softmax_num_keypoints, + freeze_backbone=freeze_backbone, + vision_backbone_mode=vision_backbone_mode, + attnres_stem_dim=attnres_stem_dim, + attnres_stage_dims=attnres_stage_dims, + attnres_stage_depths=attnres_stage_depths, + attnres_stage_heads=attnres_stage_heads, + attnres_stage_kv_heads=attnres_stage_kv_heads, + attnres_stage_window_sizes=attnres_stage_window_sizes, + attnres_dropout=attnres_dropout, + attnres_ffn_mult=attnres_ffn_mult, + attnres_eps=attnres_eps, + attnres_rope_theta=attnres_rope_theta, + ) + for _ in range(num_cameras) + ] + self.rgb_encoder = nn.ModuleList(encoders) + # 重要：output_dim 始终表示单个编码器的特征维度（与 lerobot 保持一致） + self.feature_dim = encoders[0].feature_dim + else: + # 共享编码器模式：所有摄像头共享同一个编码器 + self.rgb_encoder = _SingleRgbEncoder( + vision_backbone=vision_backbone, + pretrained_backbone_weights=pretrained_backbone_weights, + input_shape=input_shape, + crop_shape=crop_shape, + crop_is_random=crop_is_random, + use_group_norm=use_group_norm, + spatial_softmax_num_keypoints=spatial_softmax_num_keypoints, + freeze_backbone=freeze_backbone, + vision_backbone_mode=vision_backbone_mode, + attnres_stem_dim=attnres_stem_dim, + attnres_stage_dims=attnres_stage_dims, + attnres_stage_depths=attnres_stage_depths, + attnres_stage_heads=attnres_stage_heads, + attnres_stage_kv_heads=attnres_stage_kv_heads, + attnres_stage_window_sizes=attnres_stage_window_sizes, + attnres_dropout=attnres_dropout, + attnres_ffn_mult=attnres_ffn_mult, + attnres_eps=attnres_eps, + attnres_rope_theta=attnres_rope_theta, + ) + self.feature_dim = self.rgb_encoder.feature_dim + + def _ordered_camera_names(self, images) -> Tuple[str, ...]: + if self.camera_names is None: + camera_names = tuple(sorted(images.keys())) + if len(camera_names) != self.num_cameras: + raise ValueError( + f"图像输入相机数量({len(camera_names)})与 num_cameras({self.num_cameras})不一致" + ) + return camera_names + + missing = [cam_name for cam_name in self.camera_names if cam_name not in images] + if missing: + raise ValueError( + f"图像输入缺少必需相机。missing={missing}, expected={list(self.camera_names)}" + ) + return self.camera_names + + def forward(self, images): + """ + Args: + images: Dict[str, Tensor], 每个摄像头的图像 + 形状: {cam_name: (B, T, C, H, W)} + + Returns: + Tensor: (B, T, total_feature_dim) + """ + any_tensor = next(iter(images.values())) + B, T = any_tensor.shape[:2] + cam_names = self._ordered_camera_names(images) + + features_all = [] + if self.use_separate_rgb_encoder_per_camera: + # 独立编码器模式：每个摄像头使用对应的编码器 + for cam_idx, cam_name in enumerate(cam_names): + img = images[cam_name] + encoder = self.rgb_encoder[cam_idx] + features = encoder.forward_single_image(img.reshape(B * T, *img.shape[2:])) + features_all.append(features) + else: + # 共享编码器模式：所有摄像头共享同一个编码器 + for cam_name in cam_names: + img = images[cam_name] + features = self.rgb_encoder.forward_single_image(img.reshape(B * T, *img.shape[2:])) + features_all.append(features) + + if self.output_tokens_per_camera: + stacked = torch.stack(features_all, dim=1) # (B*T, num_cams, feature_dim) + return stacked.view(B, T, len(cam_names), self.feature_dim) + return torch.cat(features_all, dim=1).view(B, T, -1) + + @property + def output_dim(self): + return self.feature_dim + +if __name__ == "__main__": + print("=" * 60) + print("🚀 Testing ResNetDiffusionBackbone") + print("=" * 60) + + # Configuration + B, T = 2, 5 + C, H, W = 3, 96, 96 + crop_h, crop_w = 84, 84 + num_keypoints = 32 + feature_dim_per_cam = num_keypoints * 2 + + # Create dummy input (2 cameras) + images = { + "cam_high": torch.randn(B, T, C, H, W), + "cam_wrist": torch.randn(B, T, C, H, W) + } + num_cameras = len(images) + + # ============================================================================ + # Test 1: Shared Encoder (默认模式) + # ============================================================================ + print("\n[Test 1] Shared Encoder Mode") + print("-" * 60) + backbone_shared = ResNetDiffusionBackbone( + vision_backbone="resnet18", + pretrained_backbone_weights=None, # Speed up test + input_shape=(C, H, W), + crop_shape=(crop_h, crop_w), + crop_is_random=True, + use_group_norm=True, + spatial_softmax_num_keypoints=num_keypoints, + use_separate_rgb_encoder_per_camera=False, # 共享编码器 + ) + + print(f"✅ Shared encoder model instantiated") + print(f" Output dim per camera: {feature_dim_per_cam}") + print(f" Number of cameras: {num_cameras}") + print(f" Expected total dim: {num_cameras * feature_dim_per_cam}") + + output = backbone_shared(images) + print(f"\n🔄 Forward pass completed") + print(f" Input shapes: {[v.shape for v in images.values()]}") + print(f" Output shape: {output.shape}") + + expected_dim = num_cameras * feature_dim_per_cam + assert output.shape == (B, T, expected_dim), f"Expected shape {(B, T, expected_dim)}, got {output.shape}" + print(f"✨ Test passed!") + + # ============================================================================ + # Test 2: Separate Encoders (独立编码器模式) + # ============================================================================ + print("\n[Test 2] Separate Encoders Mode") + print("-" * 60) + backbone_separate = ResNetDiffusionBackbone( + vision_backbone="resnet18", + pretrained_backbone_weights=None, # Speed up test + input_shape=(C, H, W), + crop_shape=(crop_h, crop_w), + crop_is_random=True, + use_group_norm=True, + spatial_softmax_num_keypoints=num_keypoints, + use_separate_rgb_encoder_per_camera=True, # 独立编码器 + num_cameras=num_cameras, + ) + + print(f"✅ Separate encoders model instantiated") + print(f" Output dim per camera: {feature_dim_per_cam}") + print(f" Number of cameras: {num_cameras}") + print(f" Number of encoders: {len(backbone_separate.rgb_encoder)}") + + output = backbone_separate(images) + print(f"\n🔄 Forward pass completed") + print(f" Input shapes: {[v.shape for v in images.values()]}") + print(f" Output shape: {output.shape}") + + expected_dim = num_cameras * feature_dim_per_cam + assert output.shape == (B, T, expected_dim), f"Expected shape {(B, T, expected_dim)}, got {output.shape}" + print(f"✨ Test passed!") + + # ============================================================================ + # Test 3: Verify parameters count + # ============================================================================ + print("\n[Test 3] Parameter Count Comparison") + print("-" * 60) + shared_params = sum(p.numel() for p in backbone_shared.parameters()) + separate_params = sum(p.numel() for p in backbone_separate.parameters()) + + print(f" Shared encoder parameters: {shared_params:,}") + print(f" Separate encoders parameters: {separate_params:,}") + print(f" Ratio: {separate_params / shared_params:.2f}x") + + assert separate_params > shared_params, "Separate encoders should have more parameters" + print(f"✨ Verification passed!") + + # ============================================================================ + # Test 4: Verify independent parameters + # ============================================================================ + print("\n[Test 4] Verify Independent Parameters") + print("-" * 60) + # Check that encoders have independent parameters + encoder_0_first_param = list(backbone_separate.rgb_encoder[0].parameters())[0] + encoder_1_first_param = list(backbone_separate.rgb_encoder[1].parameters())[0] + + # Modify first encoder's parameter + with torch.no_grad(): + encoder_0_first_param += 1.0 + + # Verify they are not the same tensor + assert not torch.allclose(encoder_0_first_param, encoder_1_first_param), \ + "Encoders should have independent parameters" + + print(f"✅ Encoders have independent parameters") + print(f"✨ All tests passed!") + + print("\n" + "=" * 60) + print("🎉 All tests completed successfully!") + print("=" * 60) diff --git a/roboimi/vla/models/backbones/siglip2_diffusion_backbone.py b/roboimi/vla/models/backbones/siglip2_diffusion_backbone.py new file mode 100644 index 0000000..9a015ab --- /dev/null +++ b/roboimi/vla/models/backbones/siglip2_diffusion_backbone.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import Dict, Optional, Sequence, Tuple + +import torch +from torch import nn +from transformers import SiglipVisionModel + +from roboimi.vla.core.interfaces import VLABackbone + + +class SigLIP2DiffusionBackbone(VLABackbone): + """Shared SigLIP vision tower for multiview diffusion-policy conditioning. + + We intentionally load the checkpoint `google/siglip2-base-patch16-256` through + `SiglipVisionModel.from_pretrained(...)` so each camera can be fed as a normal + `(B, C, H, W)` image tensor and produce one pooled global feature vector. + """ + + def __init__( + self, + model_name: str = 'google/siglip2-base-patch16-256', + *, + model_name_or_path: str | None = None, + vision_model: nn.Module | None = None, + camera_names: Sequence[str] = ('r_vis', 'top', 'front'), + num_cameras: Optional[int] = None, + per_view_output_dim: int = 96, + output_dim: int | None = None, + freeze_backbone: bool = True, + dataset_image_resize_shape: Sequence[int] | None = None, + eval_image_resize_shape: Sequence[int] | None = (256, 256), + ) -> None: + super().__init__() + if model_name_or_path is not None: + model_name = model_name_or_path + if output_dim is not None: + per_view_output_dim = output_dim + + self.model_name = str(model_name) + self.camera_names = tuple(camera_names) + self.num_cameras = int(num_cameras) if num_cameras is not None else len(self.camera_names) + if len(self.camera_names) != self.num_cameras: + raise ValueError( + f'camera_names length ({len(self.camera_names)}) must match num_cameras ({self.num_cameras})' + ) + + self._output_dim = int(per_view_output_dim) + self.joint_output_dim = self._output_dim * self.num_cameras + self.freeze_backbone = bool(freeze_backbone) + self.dataset_image_resize_shape = self._normalize_resize_shape(dataset_image_resize_shape) + self.eval_image_resize_shape = self._normalize_resize_shape(eval_image_resize_shape) + + self.encoder = vision_model if vision_model is not None else SiglipVisionModel.from_pretrained(self.model_name) + hidden_size = int(getattr(self.encoder.config, 'hidden_size')) + self.view_projector = nn.Linear(hidden_size, self._output_dim) + self.projector = self.view_projector + + self.register_buffer('mean', torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32).view(1, 3, 1, 1)) + self.register_buffer('std', torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32).view(1, 3, 1, 1)) + + if self.freeze_backbone: + self._freeze_encoder() + + @staticmethod + def _normalize_resize_shape(shape: Sequence[int] | None) -> tuple[int, int] | None: + if shape is None: + return None + normalized = tuple(int(v) for v in shape) + if len(normalized) != 2: + raise ValueError(f'resize shape must contain exactly two values, got {normalized}') + return normalized + + @property + def output_dim(self) -> int: + return self._output_dim + + def _freeze_encoder(self) -> None: + self.encoder.eval() + for param in self.encoder.parameters(): + param.requires_grad = False + + def train(self, mode: bool = True): + super().train(mode) + if self.freeze_backbone: + self._freeze_encoder() + return self + + def _ordered_camera_names(self, images: Dict[str, torch.Tensor]) -> Tuple[str, ...]: + missing = [camera_name for camera_name in self.camera_names if camera_name not in images] + if missing: + raise ValueError( + f'image input missing required cameras. missing={missing}, expected={list(self.camera_names)}' + ) + return self.camera_names + + def _prepare_pixels(self, image: torch.Tensor) -> torch.Tensor: + if image.ndim != 5: + raise ValueError(f'expected image tensor shaped (B, T, C, H, W), got {tuple(image.shape)}') + pixels = image.reshape(-1, *image.shape[2:]).contiguous().float() + pixels = pixels.clamp(0.0, 1.0) + return (pixels - self.mean) / self.std + + def forward(self, images: Dict[str, torch.Tensor]) -> torch.Tensor: + camera_names = self._ordered_camera_names(images) + reference_shape = images[camera_names[0]].shape + batch_size, steps = reference_shape[:2] + per_view_features = [] + for camera_name in camera_names: + image = images[camera_name] + if image.shape != reference_shape: + raise ValueError( + f'camera {camera_name!r} shape {tuple(image.shape)} does not match {tuple(reference_shape)}' + ) + pixels = self._prepare_pixels(image) + with torch.set_grad_enabled(torch.is_grad_enabled() and not self.freeze_backbone): + encoded = self.encoder(pixel_values=pixels) + pooled = encoded.pooler_output + per_view_features.append(self.view_projector(pooled)) + features = torch.cat(per_view_features, dim=-1) + return features.view(batch_size, steps, self.joint_output_dim) + + +Siglip2DiffusionBackbone = SigLIP2DiffusionBackbone diff --git a/roboimi/vla/models/heads/__init__.py b/roboimi/vla/models/heads/__init__.py new file mode 100644 index 0000000..9e4ba5c --- /dev/null +++ b/roboimi/vla/models/heads/__init__.py @@ -0,0 +1,5 @@ +# Action Head models +from .conditional_unet1d import ConditionalUnet1D +from .transformer1d import Transformer1D + +__all__ = ["ConditionalUnet1D", "Transformer1D"] diff --git a/roboimi/vla/models/heads/attnres_transformer_components.py b/roboimi/vla/models/heads/attnres_transformer_components.py new file mode 100644 index 0000000..013eb60 --- /dev/null +++ b/roboimi/vla/models/heads/attnres_transformer_components.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) + return (x.float() * rms).to(x.dtype) * self.weight + + +class RMSNormNoWeight(nn.Module): + def __init__(self, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + + def forward(self, x: Tensor) -> Tensor: + rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) + return (x.float() * rms).to(x.dtype) + + +def precompute_rope_freqs( + dim: int, + max_seq_len: int, + theta: float = 10000.0, + device: Optional[torch.device] = None, +) -> Tensor: + if dim % 2 != 0: + raise ValueError(f'RoPE requires an even head dimension, got {dim}.') + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim)) + positions = torch.arange(max_seq_len, device=device).float() + angles = torch.outer(positions, freqs) + return torch.polar(torch.ones_like(angles), angles) + + +def apply_rope(x: Tensor, freqs: Tensor) -> Tensor: + x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs = freqs.unsqueeze(0).unsqueeze(2) + x_rotated = x_complex * freqs + return torch.view_as_real(x_rotated).reshape_as(x).to(x.dtype) + + +class GroupedQuerySelfAttention(nn.Module): + def __init__( + self, + d_model: int, + n_heads: int, + n_kv_heads: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + if d_model % n_heads != 0: + raise ValueError(f'd_model={d_model} must be divisible by n_heads={n_heads}.') + if n_heads % n_kv_heads != 0: + raise ValueError(f'n_heads={n_heads} must be divisible by n_kv_heads={n_kv_heads}.') + + self.d_model = d_model + self.n_heads = n_heads + self.n_kv_heads = n_kv_heads + self.n_kv_groups = n_heads // n_kv_heads + self.d_head = d_model // n_heads + self.attn_dropout = nn.Dropout(dropout) + self.out_dropout = nn.Dropout(dropout) + + self.w_q = nn.Linear(d_model, n_heads * self.d_head, bias=False) + self.w_k = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False) + self.w_v = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False) + self.w_o = nn.Linear(n_heads * self.d_head, d_model, bias=False) + + def forward( + self, + x: Tensor, + rope_freqs: Tensor, + mask: Optional[Tensor] = None, + ) -> Tensor: + batch_size, seq_len, _ = x.shape + + q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.d_head) + k = self.w_k(x).view(batch_size, seq_len, self.n_kv_heads, self.d_head) + v = self.w_v(x).view(batch_size, seq_len, self.n_kv_heads, self.d_head) + + q = apply_rope(q, rope_freqs) + k = apply_rope(k, rope_freqs) + + if self.n_kv_heads != self.n_heads: + k = k.unsqueeze(3).expand( + batch_size, seq_len, self.n_kv_heads, self.n_kv_groups, self.d_head + ) + k = k.reshape(batch_size, seq_len, self.n_heads, self.d_head) + v = v.unsqueeze(3).expand( + batch_size, seq_len, self.n_kv_heads, self.n_kv_groups, self.d_head + ) + v = v.reshape(batch_size, seq_len, self.n_heads, self.d_head) + + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + scale = 1.0 / math.sqrt(self.d_head) + attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale + if mask is not None: + attn_weights = attn_weights + mask + attn_weights = F.softmax(attn_weights, dim=-1) + attn_weights = self.attn_dropout(attn_weights) + + out = torch.matmul(attn_weights, v) + out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) + return self.out_dropout(self.w_o(out)) + + +class SwiGLUFFN(nn.Module): + def __init__(self, d_model: int, dropout: float = 0.0, mult: float = 2.667) -> None: + super().__init__() + raw = int(mult * d_model) + d_ff = ((raw + 7) // 8) * 8 + self.w_gate = nn.Linear(d_model, d_ff, bias=False) + self.w_up = nn.Linear(d_model, d_ff, bias=False) + self.w_down = nn.Linear(d_ff, d_model, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: Tensor) -> Tensor: + return self.dropout(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))) + + +class AttnResOperator(nn.Module): + def __init__(self, d_model: int, eps: float = 1e-6) -> None: + super().__init__() + self.pseudo_query = nn.Parameter(torch.zeros(d_model)) + self.key_norm = RMSNormNoWeight(eps=eps) + + def forward(self, sources: Tensor) -> Tensor: + keys = self.key_norm(sources) + logits = torch.einsum('d,nbtd->nbt', self.pseudo_query, keys) + weights = F.softmax(logits, dim=0) + return torch.einsum('nbt,nbtd->btd', weights, sources) + + +class AttnResSubLayer(nn.Module): + def __init__( + self, + d_model: int, + n_heads: int, + n_kv_heads: int, + dropout: float, + ffn_mult: float, + eps: float, + is_attention: bool, + ) -> None: + super().__init__() + self.norm = RMSNorm(d_model, eps=eps) + self.attn_res = AttnResOperator(d_model, eps=eps) + self.is_attention = is_attention + if self.is_attention: + self.fn = GroupedQuerySelfAttention( + d_model=d_model, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + dropout=dropout, + ) + else: + self.fn = SwiGLUFFN(d_model=d_model, dropout=dropout, mult=ffn_mult) + + def forward(self, sources: Tensor, rope_freqs: Tensor, mask: Optional[Tensor] = None) -> Tensor: + h = self.attn_res(sources) + normed = self.norm(h) + if self.is_attention: + return self.fn(normed, rope_freqs, mask) + return self.fn(normed) + + +class AttnResTransformerBackbone(nn.Module): + def __init__( + self, + d_model: int, + n_blocks: int, + n_heads: int, + n_kv_heads: int, + max_seq_len: int, + dropout: float = 0.0, + ffn_mult: float = 2.667, + eps: float = 1e-6, + rope_theta: float = 10000.0, + causal_attn: bool = False, + ) -> None: + super().__init__() + self.causal_attn = causal_attn + self.layers = nn.ModuleList() + for _ in range(n_blocks): + self.layers.append( + AttnResSubLayer( + d_model=d_model, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + dropout=dropout, + ffn_mult=ffn_mult, + eps=eps, + is_attention=True, + ) + ) + self.layers.append( + AttnResSubLayer( + d_model=d_model, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + dropout=dropout, + ffn_mult=ffn_mult, + eps=eps, + is_attention=False, + ) + ) + + rope_freqs = precompute_rope_freqs( + dim=d_model // n_heads, + max_seq_len=max_seq_len, + theta=rope_theta, + ) + self.register_buffer('rope_freqs', rope_freqs, persistent=False) + + @staticmethod + def _build_causal_mask(seq_len: int, device: torch.device) -> Tensor: + mask = torch.full((seq_len, seq_len), float('-inf'), device=device) + mask = torch.triu(mask, diagonal=1) + return mask.unsqueeze(0).unsqueeze(0) + + def forward(self, x: Tensor) -> Tensor: + seq_len = x.shape[1] + rope_freqs = self.rope_freqs[:seq_len] + mask = None + if self.causal_attn: + mask = self._build_causal_mask(seq_len, x.device) + + layer_outputs = [x] + for layer in self.layers: + sources = torch.stack(layer_outputs, dim=0) + output = layer(sources, rope_freqs, mask) + layer_outputs.append(output) + + return torch.stack(layer_outputs, dim=0).sum(dim=0) diff --git a/roboimi/vla/models/heads/conditional_unet1d.py b/roboimi/vla/models/heads/conditional_unet1d.py new file mode 100644 index 0000000..b9cc11e --- /dev/null +++ b/roboimi/vla/models/heads/conditional_unet1d.py @@ -0,0 +1,256 @@ +# Diffusion Policy Action Head 实现 +import torch +import torch.nn as nn +from typing import Dict, Optional +from diffusers import DDPMScheduler +from roboimi.vla.core.interfaces import VLAHead + +from typing import Union +import logging +import torch +import torch.nn as nn +import einops + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops.layers.torch import Rearrange +import math + + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + +class Downsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.Conv1d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + +class Upsample1d(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) + + def forward(self, x): + return self.conv(x) + +class Conv1dBlock(nn.Module): + ''' + Conv1d --> GroupNorm --> Mish + ''' + + def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): + super().__init__() + + self.block = nn.Sequential( + nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2), + # Rearrange('batch channels horizon -> batch channels 1 horizon'), + nn.GroupNorm(n_groups, out_channels), + # Rearrange('batch channels 1 horizon -> batch channels horizon'), + nn.Mish(), + ) + + def forward(self, x): + return self.block(x) + +class ConditionalResidualBlock1D(nn.Module): + def __init__(self, + in_channels, + out_channels, + cond_dim, + kernel_size=3, + n_groups=8, + cond_predict_scale=False): + super().__init__() + self.blocks = nn.ModuleList([ + Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups), + Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups), + ]) + + + + cond_channels = out_channels + if cond_predict_scale: + cond_channels = out_channels * 2 + self.cond_predict_scale = cond_predict_scale + self.out_channels = out_channels + self.cond_encoder = nn.Sequential( + nn.Mish(), + nn.Linear(cond_dim, cond_channels), + Rearrange('batch t -> batch t 1'), + ) + + # make sure dimensions compatible + self.residual_conv = nn.Conv1d(in_channels, out_channels, 1) \ + if in_channels != out_channels else nn.Identity() + + def forward(self, x, cond): + ''' + x : [ batch_size x in_channels x horizon ] + cond : [ batch_size x cond_dim] + + returns: + out : [ batch_size x out_channels x horizon ] + ''' + out = self.blocks[0](x) + embed = self.cond_encoder(cond) + if self.cond_predict_scale: + embed = embed.reshape( + embed.shape[0], 2, self.out_channels, 1) + scale = embed[:,0,...] + bias = embed[:,1,...] + out = scale * out + bias + else: + out = out + embed + out = self.blocks[1](out) + out = out + self.residual_conv(x) + return out + + +class ConditionalUnet1D(nn.Module): + def __init__(self, + input_dim, + global_cond_dim=None, + diffusion_step_embed_dim=256, + down_dims=[256,512,1024], + kernel_size=3, + n_groups=8, + cond_predict_scale=False + ): + super().__init__() + all_dims = [input_dim] + list(down_dims) + start_dim = down_dims[0] + + dsed = diffusion_step_embed_dim + diffusion_step_encoder = nn.Sequential( + SinusoidalPosEmb(dsed), + nn.Linear(dsed, dsed * 4), + nn.Mish(), + nn.Linear(dsed * 4, dsed), + ) + cond_dim = dsed + if global_cond_dim is not None: + cond_dim += global_cond_dim + + in_out = list(zip(all_dims[:-1], all_dims[1:])) + + mid_dim = all_dims[-1] + self.mid_modules = nn.ModuleList([ + ConditionalResidualBlock1D( + mid_dim, mid_dim, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups, + cond_predict_scale=cond_predict_scale + ), + ConditionalResidualBlock1D( + mid_dim, mid_dim, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups, + cond_predict_scale=cond_predict_scale + ), + ]) + + down_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (len(in_out) - 1) + down_modules.append(nn.ModuleList([ + ConditionalResidualBlock1D( + dim_in, dim_out, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups, + cond_predict_scale=cond_predict_scale), + ConditionalResidualBlock1D( + dim_out, dim_out, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups, + cond_predict_scale=cond_predict_scale), + Downsample1d(dim_out) if not is_last else nn.Identity() + ])) + + up_modules = nn.ModuleList([]) + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (len(in_out) - 1) + up_modules.append(nn.ModuleList([ + ConditionalResidualBlock1D( + dim_out*2, dim_in, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups, + cond_predict_scale=cond_predict_scale), + ConditionalResidualBlock1D( + dim_in, dim_in, cond_dim=cond_dim, + kernel_size=kernel_size, n_groups=n_groups, + cond_predict_scale=cond_predict_scale), + Upsample1d(dim_in) if not is_last else nn.Identity() + ])) + + final_conv = nn.Sequential( + Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size), + nn.Conv1d(start_dim, input_dim, 1), + ) + + self.diffusion_step_encoder = diffusion_step_encoder + self.up_modules = up_modules + self.down_modules = down_modules + self.final_conv = final_conv + + + def forward(self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + global_cond=None, + **kwargs): + """ + x: (B,T,input_dim) + timestep: (B,) or int, diffusion step + global_cond: (B,global_cond_dim) + output: (B,T,input_dim) + """ + sample = einops.rearrange(sample, 'b h t -> b t h') + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) + elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + global_feature = self.diffusion_step_encoder(timesteps) + + if global_cond is not None: + global_feature = torch.cat([ + global_feature, global_cond + ], axis=-1) + + x = sample + h = [] + for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + h.append(x) + x = downsample(x) + + for mid_module in self.mid_modules: + x = mid_module(x, global_feature) + + for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): + x = torch.cat((x, h.pop()), dim=1) + x = resnet(x, global_feature) + x = resnet2(x, global_feature) + x = upsample(x) + + x = self.final_conv(x) + + x = einops.rearrange(x, 'b t h -> b h t') + return x diff --git a/roboimi/vla/models/heads/gr00t_dit1d.py b/roboimi/vla/models/heads/gr00t_dit1d.py new file mode 100644 index 0000000..f6d6a85 --- /dev/null +++ b/roboimi/vla/models/heads/gr00t_dit1d.py @@ -0,0 +1,146 @@ +import torch +import torch.nn as nn +from types import SimpleNamespace +from typing import Optional, Union +from pathlib import Path +import importlib.util + + +def _load_gr00t_dit(): + repo_root = Path(__file__).resolve().parents[4] + dit_path = repo_root / "gr00t" / "models" / "dit.py" + spec = importlib.util.spec_from_file_location("gr00t_dit_standalone", dit_path) + if spec is None or spec.loader is None: + raise ImportError(f"Unable to load DiT from {dit_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.DiT + + +DiT = _load_gr00t_dit() + + +class Gr00tDiT1D(nn.Module): + """ + Adapter that wraps gr00t DiT with the same call signature used by VLA heads. + + Expected forward interface: + - sample: (B, T_action, input_dim) + - timestep: (B,) or scalar diffusion timestep + - cond: (B, T_obs, cond_dim) + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + horizon: int, + n_obs_steps: int, + cond_dim: int, + n_layer: int = 8, + n_head: int = 8, + n_emb: int = 256, + hidden_dim: int = 256, + mlp_ratio: int = 4, + dropout: float = 0.1, + add_action_pos_emb: bool = True, + add_cond_pos_emb: bool = True, + ): + super().__init__() + if cond_dim <= 0: + raise ValueError("Gr00tDiT1D requires cond_dim > 0.") + + self.horizon = horizon + self.n_obs_steps = n_obs_steps + + self.input_proj = nn.Linear(input_dim, n_emb) + self.cond_proj = nn.Linear(cond_dim, n_emb) + self.output_proj = nn.Linear(hidden_dim, output_dim) + + self.action_pos_emb = ( + nn.Parameter(torch.zeros(1, horizon, n_emb)) + if add_action_pos_emb + else None + ) + self.cond_pos_emb = ( + nn.Parameter(torch.zeros(1, n_obs_steps, n_emb)) + if add_cond_pos_emb + else None + ) + + args = SimpleNamespace( + embed_dim=n_emb, + nheads=n_head, + mlp_ratio=mlp_ratio, + dropout=dropout, + num_layers=n_layer, + hidden_dim=hidden_dim, + ) + self.dit = DiT(args, cross_attention_dim=n_emb) + + self._init_weights() + + def _init_weights(self): + if self.action_pos_emb is not None: + nn.init.normal_(self.action_pos_emb, mean=0.0, std=0.02) + if self.cond_pos_emb is not None: + nn.init.normal_(self.cond_pos_emb, mean=0.0, std=0.02) + + def _normalize_timesteps( + self, + timestep: Union[torch.Tensor, float, int], + batch_size: int, + device: torch.device, + ) -> torch.Tensor: + if not torch.is_tensor(timestep): + timesteps = torch.tensor([timestep], device=device) + else: + timesteps = timestep.to(device) + + if timesteps.ndim == 0: + timesteps = timesteps[None] + if timesteps.shape[0] != batch_size: + timesteps = timesteps.expand(batch_size) + + return timesteps.long() + + def forward( + self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + cond: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + if cond is None: + raise ValueError("`cond` is required for Gr00tDiT1D forward.") + + bsz, t_act, _ = sample.shape + if t_act > self.horizon: + raise ValueError( + f"sample length {t_act} exceeds configured horizon {self.horizon}" + ) + + hidden_states = self.input_proj(sample) + if self.action_pos_emb is not None: + hidden_states = hidden_states + self.action_pos_emb[:, :t_act, :] + + encoder_hidden_states = self.cond_proj(cond) + if self.cond_pos_emb is not None: + t_obs = encoder_hidden_states.shape[1] + if t_obs > self.n_obs_steps: + raise ValueError( + f"cond length {t_obs} exceeds configured n_obs_steps {self.n_obs_steps}" + ) + encoder_hidden_states = ( + encoder_hidden_states + self.cond_pos_emb[:, :t_obs, :] + ) + + timesteps = self._normalize_timesteps( + timestep, batch_size=bsz, device=sample.device + ) + dit_output = self.dit( + hidden_states=hidden_states, + timestep=timesteps, + encoder_hidden_states=encoder_hidden_states, + ) + return self.output_proj(dit_output) diff --git a/roboimi/vla/models/heads/imf_transformer1d.py b/roboimi/vla/models/heads/imf_transformer1d.py new file mode 100644 index 0000000..ae605c1 --- /dev/null +++ b/roboimi/vla/models/heads/imf_transformer1d.py @@ -0,0 +1,379 @@ +"""Local IMF-AttnRes transformer head aligned with diffusion_policy@185ed659.""" + +from __future__ import annotations + +import logging +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +from .attnres_transformer_components import ( + AttnResOperator, + AttnResSubLayer, + AttnResTransformerBackbone, + GroupedQuerySelfAttention, + RMSNorm, + RMSNormNoWeight, + SwiGLUFFN, +) +from .transformer1d import ModuleAttrMixin, SinusoidalPosEmb + +logger = logging.getLogger(__name__) + + +class IMFTransformer1D(ModuleAttrMixin): + def __init__( + self, + input_dim: int, + output_dim: int, + horizon: int, + n_obs_steps: Optional[int] = None, + cond_dim: int = 0, + n_layer: int = 12, + n_head: int = 1, + n_emb: int = 768, + p_drop_emb: float = 0.1, + p_drop_attn: float = 0.1, + causal_attn: bool = False, + time_as_cond: bool = True, + obs_as_cond: bool = False, + n_cond_layers: int = 0, + backbone_type: str = 'attnres_full', + n_kv_head: int = 1, + attn_res_ffn_mult: float = 2.667, + attn_res_eps: float = 1e-6, + attn_res_rope_theta: float = 10000.0, + ) -> None: + super().__init__() + + if n_head != 1: + raise AssertionError('IMFTransformer1D currently supports single-head attention only.') + if n_obs_steps is None: + n_obs_steps = horizon + + self.backbone_type = backbone_type + + T = horizon + T_cond = 2 + if not time_as_cond: + T += 2 + T_cond -= 2 + obs_as_cond = cond_dim > 0 + if obs_as_cond: + assert time_as_cond + T_cond += n_obs_steps + + self.input_emb = nn.Linear(input_dim, n_emb) + self.drop = nn.Dropout(p_drop_emb) + self.time_emb = SinusoidalPosEmb(n_emb) + self.cond_obs_emb = nn.Linear(cond_dim, n_emb) if obs_as_cond else None + self.time_token_proj = None + self.cond_pos_emb = None + self.pos_emb = None + self.encoder = None + self.decoder = None + self.attnres_backbone = None + encoder_only = False + + if backbone_type == 'attnres_full': + if not time_as_cond: + raise ValueError('attnres_full backbone requires time_as_cond=True.') + if n_cond_layers != 0: + raise ValueError('attnres_full backbone does not support n_cond_layers > 0.') + + self.time_token_proj = nn.Linear(n_emb, n_emb) + self.attnres_backbone = AttnResTransformerBackbone( + d_model=n_emb, + n_blocks=n_layer, + n_heads=n_head, + n_kv_heads=n_kv_head, + max_seq_len=T + T_cond, + dropout=p_drop_attn, + ffn_mult=attn_res_ffn_mult, + eps=attn_res_eps, + rope_theta=attn_res_rope_theta, + causal_attn=causal_attn, + ) + self.ln_f = RMSNorm(n_emb, eps=attn_res_eps) + else: + self.pos_emb = nn.Parameter(torch.zeros(1, T, n_emb)) + if T_cond > 0: + self.cond_pos_emb = nn.Parameter(torch.zeros(1, T_cond, n_emb)) + if n_cond_layers > 0: + encoder_layer = nn.TransformerEncoderLayer( + d_model=n_emb, + nhead=n_head, + dim_feedforward=4 * n_emb, + dropout=p_drop_attn, + activation='gelu', + batch_first=True, + norm_first=True, + ) + self.encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=n_cond_layers, + ) + else: + self.encoder = nn.Sequential( + nn.Linear(n_emb, 4 * n_emb), + nn.Mish(), + nn.Linear(4 * n_emb, n_emb), + ) + + decoder_layer = nn.TransformerDecoderLayer( + d_model=n_emb, + nhead=n_head, + dim_feedforward=4 * n_emb, + dropout=p_drop_attn, + activation='gelu', + batch_first=True, + norm_first=True, + ) + self.decoder = nn.TransformerDecoder( + decoder_layer=decoder_layer, + num_layers=n_layer, + ) + else: + encoder_only = True + encoder_layer = nn.TransformerEncoderLayer( + d_model=n_emb, + nhead=n_head, + dim_feedforward=4 * n_emb, + dropout=p_drop_attn, + activation='gelu', + batch_first=True, + norm_first=True, + ) + self.encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=n_layer, + ) + + self.ln_f = nn.LayerNorm(n_emb) + + if causal_attn and backbone_type != 'attnres_full': + sz = T + mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + self.register_buffer('mask', mask) + + if time_as_cond and obs_as_cond: + S = T_cond + t_idx, s_idx = torch.meshgrid( + torch.arange(T), + torch.arange(S), + indexing='ij', + ) + mask = t_idx >= (s_idx - 2) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + self.register_buffer('memory_mask', mask) + else: + self.memory_mask = None + else: + self.mask = None + self.memory_mask = None + + self.head = nn.Linear(n_emb, output_dim) + + self.T = T + self.T_cond = T_cond + self.horizon = horizon + self.time_as_cond = time_as_cond + self.obs_as_cond = obs_as_cond + self.encoder_only = encoder_only + + self.apply(self._init_weights) + logger.info('number of parameters: %e', sum(p.numel() for p in self.parameters())) + + def _init_weights(self, module): + ignore_types = ( + nn.Dropout, + SinusoidalPosEmb, + nn.TransformerEncoderLayer, + nn.TransformerDecoderLayer, + nn.TransformerEncoder, + nn.TransformerDecoder, + nn.ModuleList, + nn.Mish, + nn.Sequential, + AttnResTransformerBackbone, + AttnResSubLayer, + GroupedQuerySelfAttention, + SwiGLUFFN, + RMSNormNoWeight, + ) + if isinstance(module, (nn.Linear, nn.Embedding)): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.MultiheadAttention): + for name in ('in_proj_weight', 'q_proj_weight', 'k_proj_weight', 'v_proj_weight'): + weight = getattr(module, name) + if weight is not None: + torch.nn.init.normal_(weight, mean=0.0, std=0.02) + + for name in ('in_proj_bias', 'bias_k', 'bias_v'): + bias = getattr(module, name) + if bias is not None: + torch.nn.init.zeros_(bias) + elif isinstance(module, (nn.LayerNorm, RMSNorm)): + if getattr(module, 'bias', None) is not None: + torch.nn.init.zeros_(module.bias) + torch.nn.init.ones_(module.weight) + elif isinstance(module, AttnResOperator): + torch.nn.init.zeros_(module.pseudo_query) + elif isinstance(module, IMFTransformer1D): + if module.pos_emb is not None: + torch.nn.init.normal_(module.pos_emb, mean=0.0, std=0.02) + if module.cond_pos_emb is not None: + torch.nn.init.normal_(module.cond_pos_emb, mean=0.0, std=0.02) + elif isinstance(module, ignore_types): + pass + else: + raise RuntimeError(f'Unaccounted module {module}') + + def get_optim_groups(self, weight_decay: float = 1e-3): + decay = set() + no_decay = set() + whitelist_weight_modules = (torch.nn.Linear, torch.nn.MultiheadAttention) + blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding, RMSNorm) + for mn, m in self.named_modules(): + for pn, _ in m.named_parameters(recurse=False): + fpn = f'{mn}.{pn}' if mn else pn + + if pn.endswith('bias'): + no_decay.add(fpn) + elif pn.startswith('bias'): + no_decay.add(fpn) + elif pn == 'pseudo_query': + no_decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules): + decay.add(fpn) + elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules): + no_decay.add(fpn) + + if self.pos_emb is not None: + no_decay.add('pos_emb') + no_decay.add('_dummy_variable') + if self.cond_pos_emb is not None: + no_decay.add('cond_pos_emb') + + param_dict = {pn: p for pn, p in self.named_parameters()} + inter_params = decay & no_decay + union_params = decay | no_decay + assert len(inter_params) == 0, f'parameters {inter_params} made it into both decay/no_decay sets!' + assert len(param_dict.keys() - union_params) == 0, ( + f'parameters {param_dict.keys() - union_params} were not separated into either decay/no_decay sets!' + ) + + return [ + { + 'params': [param_dict[pn] for pn in sorted(list(decay))], + 'weight_decay': weight_decay, + }, + { + 'params': [param_dict[pn] for pn in sorted(list(no_decay))], + 'weight_decay': 0.0, + }, + ] + + def configure_optimizers( + self, + learning_rate: float = 1e-4, + weight_decay: float = 1e-3, + betas: Tuple[float, float] = (0.9, 0.95), + ): + optim_groups = self.get_optim_groups(weight_decay=weight_decay) + return torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas) + + def _prepare_time_input(self, value: Union[torch.Tensor, float, int], sample: torch.Tensor) -> torch.Tensor: + if not torch.is_tensor(value): + value = torch.tensor([value], dtype=sample.dtype, device=sample.device) + elif value.ndim == 0: + value = value[None].to(device=sample.device, dtype=sample.dtype) + else: + value = value.to(device=sample.device, dtype=sample.dtype) + return value.expand(sample.shape[0]) + + def _forward_attnres_full( + self, + sample: torch.Tensor, + r: torch.Tensor, + t: torch.Tensor, + cond: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + sample_tokens = self.input_emb(sample) + token_parts = [ + self.time_token_proj(self.time_emb(r)).unsqueeze(1), + self.time_token_proj(self.time_emb(t)).unsqueeze(1), + ] + if self.obs_as_cond: + if cond is None: + raise ValueError('cond is required when obs_as_cond=True for attnres_full backbone.') + token_parts.append(self.cond_obs_emb(cond)) + token_parts.append(sample_tokens) + x = torch.cat(token_parts, dim=1) + x = self.drop(x) + x = self.attnres_backbone(x) + x = x[:, -sample_tokens.shape[1]:, :] + return x + + def _forward_vanilla( + self, + sample: torch.Tensor, + r: torch.Tensor, + t: torch.Tensor, + cond: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + r_emb = self.time_emb(r).unsqueeze(1) + t_emb = self.time_emb(t).unsqueeze(1) + input_emb = self.input_emb(sample) + + if self.encoder_only: + token_embeddings = torch.cat([r_emb, t_emb, input_emb], dim=1) + token_count = token_embeddings.shape[1] + position_embeddings = self.pos_emb[:, :token_count, :] + x = self.drop(token_embeddings + position_embeddings) + x = self.encoder(src=x, mask=self.mask) + x = x[:, 2:, :] + else: + cond_embeddings = torch.cat([r_emb, t_emb], dim=1) + if self.obs_as_cond: + cond_embeddings = torch.cat([cond_embeddings, self.cond_obs_emb(cond)], dim=1) + token_count = cond_embeddings.shape[1] + position_embeddings = self.cond_pos_emb[:, :token_count, :] + x = self.drop(cond_embeddings + position_embeddings) + x = self.encoder(x) + memory = x + + token_embeddings = input_emb + token_count = token_embeddings.shape[1] + position_embeddings = self.pos_emb[:, :token_count, :] + x = self.drop(token_embeddings + position_embeddings) + x = self.decoder( + tgt=x, + memory=memory, + tgt_mask=self.mask, + memory_mask=self.memory_mask, + ) + return x + + def forward( + self, + sample: torch.Tensor, + r: Union[torch.Tensor, float, int], + t: Union[torch.Tensor, float, int], + cond: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + r = self._prepare_time_input(r, sample) + t = self._prepare_time_input(t, sample) + + if self.backbone_type == 'attnres_full': + x = self._forward_attnres_full(sample, r, t, cond=cond) + else: + x = self._forward_vanilla(sample, r, t, cond=cond) + + x = self.ln_f(x) + x = self.head(x) + return x diff --git a/roboimi/vla/models/heads/transformer1d.py b/roboimi/vla/models/heads/transformer1d.py new file mode 100644 index 0000000..2b0752a --- /dev/null +++ b/roboimi/vla/models/heads/transformer1d.py @@ -0,0 +1,337 @@ +"""Transformer-based diffusion head aligned with diffusion_policy's TransformerForDiffusion.""" + +from __future__ import annotations + +import logging +import math +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +class ModuleAttrMixin(nn.Module): + """Minimal local copy of diffusion_policy's ModuleAttrMixin for state-dict parity.""" + + def __init__(self) -> None: + super().__init__() + self._dummy_variable = nn.Parameter() + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class Transformer1D(ModuleAttrMixin): + def __init__( + self, + input_dim: int, + output_dim: int, + horizon: int, + n_obs_steps: Optional[int] = None, + cond_dim: int = 0, + n_layer: int = 8, + n_head: int = 8, + n_emb: int = 256, + p_drop_emb: float = 0.1, + p_drop_attn: float = 0.1, + causal_attn: bool = False, + time_as_cond: bool = True, + obs_as_cond: bool = False, + n_cond_layers: int = 0, + ) -> None: + super().__init__() + + if n_obs_steps is None: + n_obs_steps = horizon + + T = horizon + T_cond = 1 + if not time_as_cond: + T += 1 + T_cond -= 1 + obs_as_cond = cond_dim > 0 + if obs_as_cond: + assert time_as_cond + T_cond += n_obs_steps + + self.input_emb = nn.Linear(input_dim, n_emb) + self.pos_emb = nn.Parameter(torch.zeros(1, T, n_emb)) + self.drop = nn.Dropout(p_drop_emb) + + self.time_emb = SinusoidalPosEmb(n_emb) + self.cond_obs_emb = None + if obs_as_cond: + self.cond_obs_emb = nn.Linear(cond_dim, n_emb) + + self.cond_pos_emb = None + self.encoder = None + self.decoder = None + encoder_only = False + + if T_cond > 0: + self.cond_pos_emb = nn.Parameter(torch.zeros(1, T_cond, n_emb)) + if n_cond_layers > 0: + encoder_layer = nn.TransformerEncoderLayer( + d_model=n_emb, + nhead=n_head, + dim_feedforward=4 * n_emb, + dropout=p_drop_attn, + activation='gelu', + batch_first=True, + norm_first=True, + ) + self.encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=n_cond_layers, + ) + else: + self.encoder = nn.Sequential( + nn.Linear(n_emb, 4 * n_emb), + nn.Mish(), + nn.Linear(4 * n_emb, n_emb), + ) + + decoder_layer = nn.TransformerDecoderLayer( + d_model=n_emb, + nhead=n_head, + dim_feedforward=4 * n_emb, + dropout=p_drop_attn, + activation='gelu', + batch_first=True, + norm_first=True, + ) + self.decoder = nn.TransformerDecoder( + decoder_layer=decoder_layer, + num_layers=n_layer, + ) + else: + encoder_only = True + encoder_layer = nn.TransformerEncoderLayer( + d_model=n_emb, + nhead=n_head, + dim_feedforward=4 * n_emb, + dropout=p_drop_attn, + activation='gelu', + batch_first=True, + norm_first=True, + ) + self.encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=n_layer, + ) + + if causal_attn: + sz = T + mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + self.register_buffer('mask', mask) + + if time_as_cond and obs_as_cond: + S = T_cond + t, s = torch.meshgrid(torch.arange(T), torch.arange(S), indexing='ij') + mask = t >= (s - 1) + mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) + self.register_buffer('memory_mask', mask) + else: + self.memory_mask = None + else: + self.mask = None + self.memory_mask = None + + self.ln_f = nn.LayerNorm(n_emb) + self.head = nn.Linear(n_emb, output_dim) + + self.T = T + self.T_cond = T_cond + self.horizon = horizon + self.time_as_cond = time_as_cond + self.obs_as_cond = obs_as_cond + self.encoder_only = encoder_only + + self.apply(self._init_weights) + logger.info('number of parameters: %e', sum(p.numel() for p in self.parameters())) + + def _init_weights(self, module): + ignore_types = ( + nn.Dropout, + SinusoidalPosEmb, + nn.TransformerEncoderLayer, + nn.TransformerDecoderLayer, + nn.TransformerEncoder, + nn.TransformerDecoder, + nn.ModuleList, + nn.Mish, + nn.Sequential, + ) + if isinstance(module, (nn.Linear, nn.Embedding)): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.MultiheadAttention): + for name in ('in_proj_weight', 'q_proj_weight', 'k_proj_weight', 'v_proj_weight'): + weight = getattr(module, name) + if weight is not None: + torch.nn.init.normal_(weight, mean=0.0, std=0.02) + + for name in ('in_proj_bias', 'bias_k', 'bias_v'): + bias = getattr(module, name) + if bias is not None: + torch.nn.init.zeros_(bias) + elif isinstance(module, nn.LayerNorm): + torch.nn.init.zeros_(module.bias) + torch.nn.init.ones_(module.weight) + elif isinstance(module, Transformer1D): + torch.nn.init.normal_(module.pos_emb, mean=0.0, std=0.02) + if module.cond_obs_emb is not None: + torch.nn.init.normal_(module.cond_pos_emb, mean=0.0, std=0.02) + elif isinstance(module, ignore_types): + pass + else: + raise RuntimeError(f'Unaccounted module {module}') + + def get_optim_groups(self, weight_decay: float = 1e-3): + decay = set() + no_decay = set() + whitelist_weight_modules = (torch.nn.Linear, torch.nn.MultiheadAttention) + blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) + + for module_name, module in self.named_modules(): + for param_name, _ in module.named_parameters(): + full_param_name = f'{module_name}.{param_name}' if module_name else param_name + + if param_name.endswith('bias'): + no_decay.add(full_param_name) + elif param_name.startswith('bias'): + no_decay.add(full_param_name) + elif param_name.endswith('weight') and isinstance(module, whitelist_weight_modules): + decay.add(full_param_name) + elif param_name.endswith('weight') and isinstance(module, blacklist_weight_modules): + no_decay.add(full_param_name) + + no_decay.add('pos_emb') + no_decay.add('_dummy_variable') + if self.cond_pos_emb is not None: + no_decay.add('cond_pos_emb') + + param_dict = {name: param for name, param in self.named_parameters()} + inter_params = decay & no_decay + union_params = decay | no_decay + assert len(inter_params) == 0, f'parameters {inter_params} made it into both decay/no_decay sets!' + assert len(param_dict.keys() - union_params) == 0, ( + f'parameters {param_dict.keys() - union_params} were not separated into either decay/no_decay sets!' + ) + + return [ + { + 'params': [param_dict[name] for name in sorted(decay)], + 'weight_decay': weight_decay, + }, + { + 'params': [param_dict[name] for name in sorted(no_decay)], + 'weight_decay': 0.0, + }, + ] + + def configure_optimizers( + self, + learning_rate: float = 1e-4, + weight_decay: float = 1e-3, + betas: Tuple[float, float] = (0.9, 0.95), + ): + optim_groups = self.get_optim_groups(weight_decay=weight_decay) + return torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas) + + def forward( + self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + cond: Optional[torch.Tensor] = None, + **kwargs, + ): + timesteps = timestep + if not torch.is_tensor(timesteps): + timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device) + elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + timesteps = timesteps.expand(sample.shape[0]) + time_emb = self.time_emb(timesteps).unsqueeze(1) + + input_emb = self.input_emb(sample) + + if self.encoder_only: + token_embeddings = torch.cat([time_emb, input_emb], dim=1) + t = token_embeddings.shape[1] + position_embeddings = self.pos_emb[:, :t, :] + x = self.drop(token_embeddings + position_embeddings) + x = self.encoder(src=x, mask=self.mask) + x = x[:, 1:, :] + else: + cond_embeddings = time_emb + if self.obs_as_cond: + cond_obs_emb = self.cond_obs_emb(cond) + cond_embeddings = torch.cat([cond_embeddings, cond_obs_emb], dim=1) + tc = cond_embeddings.shape[1] + position_embeddings = self.cond_pos_emb[:, :tc, :] + x = self.drop(cond_embeddings + position_embeddings) + memory = self.encoder(x) + + token_embeddings = input_emb + t = token_embeddings.shape[1] + position_embeddings = self.pos_emb[:, :t, :] + x = self.drop(token_embeddings + position_embeddings) + x = self.decoder( + tgt=x, + memory=memory, + tgt_mask=self.mask, + memory_mask=self.memory_mask, + ) + + x = self.ln_f(x) + x = self.head(x) + return x + + +def create_transformer1d( + input_dim: int, + output_dim: int, + horizon: int, + n_obs_steps: int, + cond_dim: int, + n_layer: int = 8, + n_head: int = 8, + n_emb: int = 256, + **kwargs, +) -> Transformer1D: + return Transformer1D( + input_dim=input_dim, + output_dim=output_dim, + horizon=horizon, + n_obs_steps=n_obs_steps, + cond_dim=cond_dim, + n_layer=n_layer, + n_head=n_head, + n_emb=n_emb, + **kwargs, + ) diff --git a/roboimi/vla/models/normalization.py b/roboimi/vla/models/normalization.py new file mode 100644 index 0000000..cb5adef --- /dev/null +++ b/roboimi/vla/models/normalization.py @@ -0,0 +1,126 @@ +""" +归一化模块 - 统一训练和推理的归一化逻辑 + +支持两种归一化方式: +1. Gaussian (z-score): (x - mean) / std +2. MinMax: 2 * (x - min) / (max - min) - 1 -> [-1, 1] +""" + +import torch +import torch.nn as nn +from typing import Optional, Dict, Literal + + +class NormalizationModule(nn.Module): + """ + 统一的归一化模块 + 用于在 Agent 内部对 qpos 和 action 进行归一化/反归一化 + """ + + def __init__( + self, + stats: Optional[Dict] = None, + normalization_type: Literal['gaussian', 'min_max'] = None, + ): + """ + Args: + stats: 数据集统计信息字典，格式: + { + 'qpos_mean': [...], + 'qpos_std': [...], + 'qpos_min': [...], # 仅 min_max 需要 + 'qpos_max': [...], # 仅 min_max 需要 + 'action_mean': [...], + 'action_std': [...], + 'action_min': [...], # 仅 min_max 需要 + 'action_max': [...], # 仅 min_max 需要 + } + normalization_type: 归一化类型 ('gaussian' 或 'min_max') + """ + super().__init__() + + self.normalization_type = normalization_type + self.enabled = stats is not None + + if self.enabled: + if self.normalization_type == 'gaussian': + self.register_buffer('qpos_mean', torch.tensor(stats['qpos_mean'], dtype=torch.float32)) + self.register_buffer('qpos_std', torch.tensor(stats['qpos_std'], dtype=torch.float32)) + self.register_buffer('action_mean', torch.tensor(stats['action_mean'], dtype=torch.float32)) + self.register_buffer('action_std', torch.tensor(stats['action_std'], dtype=torch.float32)) + + elif self.normalization_type == 'min_max': + self.register_buffer('qpos_min', torch.tensor(stats['qpos_min'], dtype=torch.float32)) + self.register_buffer('qpos_max', torch.tensor(stats['qpos_max'], dtype=torch.float32)) + self.register_buffer('action_min', torch.tensor(stats['action_min'], dtype=torch.float32)) + self.register_buffer('action_max', torch.tensor(stats['action_max'], dtype=torch.float32)) + + def normalize_qpos(self, qpos: torch.Tensor) -> torch.Tensor: + """归一化 qpos""" + if not self.enabled: + return qpos + + if self.normalization_type == 'gaussian': + return (qpos - self.qpos_mean) / self.qpos_std + elif self.normalization_type == 'min_max': + return 2 * (qpos - self.qpos_min) / (self.qpos_max - self.qpos_min) - 1 + else: + raise ValueError(f"Unknown normalization type: {self.normalization_type}") + + def denormalize_qpos(self, qpos: torch.Tensor) -> torch.Tensor: + """反归一化 qpos""" + if not self.enabled: + return qpos + + if self.normalization_type == 'gaussian': + return qpos * self.qpos_std + self.qpos_mean + elif self.normalization_type == 'min_max': + return (qpos + 1) / 2 * (self.qpos_max - self.qpos_min) + self.qpos_min + else: + raise ValueError(f"Unknown normalization type: {self.normalization_type}") + + def normalize_action(self, action: torch.Tensor) -> torch.Tensor: + """归一化 action""" + if not self.enabled: + return action + + if self.normalization_type == 'gaussian': + return (action - self.action_mean) / self.action_std + elif self.normalization_type == 'min_max': + return 2 * (action - self.action_min) / (self.action_max - self.action_min) - 1 + else: + raise ValueError(f"Unknown normalization type: {self.normalization_type}") + + def denormalize_action(self, action: torch.Tensor) -> torch.Tensor: + """反归一化 action""" + if not self.enabled: + return action + + if self.normalization_type == 'gaussian': + return action * self.action_std + self.action_mean + elif self.normalization_type == 'min_max': + return (action + 1) / 2 * (self.action_max - self.action_min) + self.action_min + else: + raise ValueError(f"Unknown normalization type: {self.normalization_type}") + + def get_stats(self) -> Optional[Dict]: + """导出统计信息（用于保存到 checkpoint）""" + if not self.enabled: + return None + + stats = { + 'normalization_type': self.normalization_type, + } + + if self.normalization_type == 'gaussian': + stats['qpos_mean'] = self.qpos_mean.cpu().tolist() + stats['qpos_std'] = self.qpos_std.cpu().tolist() + stats['action_mean'] = self.action_mean.cpu().tolist() + stats['action_std'] = self.action_std.cpu().tolist() + elif self.normalization_type == 'min_max': + stats['qpos_min'] = self.qpos_min.cpu().tolist() + stats['qpos_max'] = self.qpos_max.cpu().tolist() + stats['action_min'] = self.action_min.cpu().tolist() + stats['action_max'] = self.action_max.cpu().tolist() + + return stats diff --git a/roboimi/vla/modules/encoders.py b/roboimi/vla/modules/encoders.py new file mode 100644 index 0000000..0fa0970 --- /dev/null +++ b/roboimi/vla/modules/encoders.py @@ -0,0 +1,18 @@ +from torch import nn + +class IdentityStateEncoder(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, state): + return state + + +class IdentityActionEncoder(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, action): + return action \ No newline at end of file diff --git a/roboimi/vla/modules/projectors.py b/roboimi/vla/modules/projectors.py new file mode 100644 index 0000000..50e61a1 --- /dev/null +++ b/roboimi/vla/modules/projectors.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +import torch +from torch import nn + + +class LinearConditionProjector(nn.Module): + """Projects per-step visual+state conditioning to the head conditioning width.""" + + def __init__(self, input_dim: int, output_dim: int, bias: bool = True): + super().__init__() + self.input_dim = int(input_dim) + self.output_dim = int(output_dim) + self.linear = nn.Linear(self.input_dim, self.output_dim, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) diff --git a/roboimi/vla/scripts/calculate_stats.py b/roboimi/vla/scripts/calculate_stats.py new file mode 100644 index 0000000..072f4bf --- /dev/null +++ b/roboimi/vla/scripts/calculate_stats.py @@ -0,0 +1,114 @@ +import argparse +import glob +import os +import pickle +from pathlib import Path + +import h5py +import numpy as np + + +DEFAULT_DATASET_DIR = str( + Path(__file__).resolve().parents[2] / "demos" / "dataset" / "sim_transfer" +) + +def get_data_stats(dataset_dir): + """ + 计算 action 和 qpos 的 Min, Max, Mean, Std + + 输出扁平化结构（与 NormalizationModule 期望一致）： + { + 'action_mean': [...], + 'action_std': [...], + 'action_min': [...], + 'action_max': [...], + 'qpos_mean': [...], + 'qpos_std': [...], + 'qpos_min': [...], + 'qpos_max': [...], + } + """ + files = sorted(glob.glob(os.path.join(dataset_dir, 'episode_*.hdf5'))) + print(f"Found {len(files)} episodes in {dataset_dir}") + + if not files: + raise ValueError( + f"No episode_*.hdf5 files found in dataset_dir: {dataset_dir}" + ) + + all_actions = [] + all_qpos = [] + + print("Reading data...") + for file_path in files: + with h5py.File(file_path, 'r') as f: + action = f['action'][:] + qpos = f['observations']['qpos'][:] + all_actions.append(action) + all_qpos.append(qpos) + + # 拼接所有数据 + all_actions = np.concatenate(all_actions, axis=0) + all_qpos = np.concatenate(all_qpos, axis=0) + + print(f"Total steps: {all_actions.shape[0]}") + + # --- 核心计算部分 --- + # 计算统计量 + action_mean = np.mean(all_actions, axis=0) + action_std = np.std(all_actions, axis=0) + action_min = np.min(all_actions, axis=0) + action_max = np.max(all_actions, axis=0) + + qpos_mean = np.mean(all_qpos, axis=0) + qpos_std = np.std(all_qpos, axis=0) + qpos_min = np.min(all_qpos, axis=0) + qpos_max = np.max(all_qpos, axis=0) + + # 修正标准差（防止除以 0） + eps = 1e-8 + action_std_corrected = np.where(action_std < eps, eps, action_std) + qpos_std_corrected = np.where(qpos_std < eps, eps, qpos_std) + + # 转换为扁平化结构（与 NormalizationModule 期望一致） + stats_flat = { + 'action_mean': action_mean, + 'action_std': action_std_corrected, + 'action_min': action_min, + 'action_max': action_max, + 'qpos_mean': qpos_mean, + 'qpos_std': qpos_std_corrected, + 'qpos_min': qpos_min, + 'qpos_max': qpos_max + } + return stats_flat + + +def write_dataset_stats(dataset_dir): + output_path = os.path.join(dataset_dir, "dataset_stats.pkl") + stats_flat = get_data_stats(dataset_dir) + + # 打印检查 + print("\n--- Stats Computed ---") + print(f"Action Mean shape: {stats_flat['action_mean'].shape}") + print(f"Action Std shape: {stats_flat['action_std'].shape}") + + with open(output_path, 'wb') as f: + pickle.dump(stats_flat, f) + print(f"\nStats saved to {output_path}") + + return output_path + + +def main(argv=None): + parser = argparse.ArgumentParser(description="Calculate dataset statistics.") + parser.add_argument( + "--dataset_dir", + default=DEFAULT_DATASET_DIR, + help="Directory containing episode_*.hdf5 files.", + ) + args = parser.parse_args(argv) + write_dataset_stats(args.dataset_dir) + +if __name__ == "__main__": + main() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_air_insert_env.py b/tests/test_air_insert_env.py new file mode 100644 index 0000000..2c725c5 --- /dev/null +++ b/tests/test_air_insert_env.py @@ -0,0 +1,510 @@ +import importlib +import inspect +import pathlib +import unittest +from unittest import mock +import xml.etree.ElementTree as ET + +import numpy as np + +from roboimi.envs.double_pos_ctrl_env import make_sim_env +from roboimi.utils import act_ex_utils +from roboimi.utils.constants import SIM_TASK_CONFIGS + + +TASK_NAME = "sim_air_insert_socket_peg" + + +class AirInsertTaskRegistrationTest(unittest.TestCase): + def test_sim_task_configs_registers_air_insert_socket_peg(self): + self.assertIn(TASK_NAME, SIM_TASK_CONFIGS) + self.assertNotIn("sim_air_insert_ring_bar", SIM_TASK_CONFIGS) + self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"], 750) + self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["camera_names"], ["l_vis", "r_vis", "front"]) + self.assertTrue(SIM_TASK_CONFIGS[TASK_NAME]["dataset_dir"].endswith("/sim_air_insert_socket_peg")) + + def test_sample_air_insert_socket_peg_state_returns_explicit_named_mapping(self): + sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None) + self.assertIsNotNone( + sampler, + "Expected roboimi.utils.act_ex_utils.sample_air_insert_socket_peg_state()", + ) + self.assertFalse( + hasattr(act_ex_utils, "sample_air_insert_ring_bar_state"), + "air insert sampler should use socket/peg naming after the task rename", + ) + + task_state = sampler() + + self.assertEqual( + list(task_state.keys()), + ["socket_pos", "socket_quat", "peg_pos", "peg_quat"], + ) + self.assertEqual(task_state["socket_pos"].shape, (3,)) + self.assertEqual(task_state["socket_quat"].shape, (4,)) + self.assertEqual(task_state["peg_pos"].shape, (3,)) + self.assertEqual(task_state["peg_quat"].shape, (4,)) + + def test_sample_air_insert_socket_peg_state_uses_fixed_quats_and_left_right_planar_ranges(self): + sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None) + self.assertIsNotNone(sampler) + + task_state = sampler() + + np.testing.assert_array_equal(task_state["socket_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)) + np.testing.assert_array_equal(task_state["peg_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)) + self.assertGreaterEqual(task_state["socket_pos"][0], -0.20) + self.assertLessEqual(task_state["socket_pos"][0], -0.05) + self.assertGreaterEqual(task_state["socket_pos"][1], 0.70) + self.assertLessEqual(task_state["socket_pos"][1], 1.00) + self.assertAlmostEqual(float(task_state["socket_pos"][2]), 0.472) + self.assertGreaterEqual(task_state["peg_pos"][0], 0.05) + self.assertLessEqual(task_state["peg_pos"][0], 0.20) + self.assertGreaterEqual(task_state["peg_pos"][1], 0.70) + self.assertLessEqual(task_state["peg_pos"][1], 1.00) + self.assertAlmostEqual(float(task_state["peg_pos"][2]), 0.46) + + def test_make_sim_env_dispatches_air_insert_socket_peg_headless(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + air_insert_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None) + self.assertIsNotNone(air_insert_cls) + + diana_med = importlib.import_module("roboimi.assets.robots.diana_med") + socket_peg_robot_cls = getattr(diana_med, "BiDianaMedSocketPeg", None) + self.assertIsNotNone( + socket_peg_robot_cls, + "Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg", + ) + + fake_env = object() + with mock.patch.object( + diana_med, + "BiDianaMedSocketPeg", + return_value="robot", + ), mock.patch.object( + air_insert_env, + "DualDianaMed_Air_Insert", + return_value=fake_env, + ) as env_cls: + env = make_sim_env(TASK_NAME, headless=True) + + self.assertIs(env, fake_env) + env_cls.assert_called_once_with( + robot="robot", + is_render=False, + control_freq=30, + is_interpolate=True, + cam_view="front", + ) + + def test_diana_table_scene_exposes_only_top_and_front_scene_cameras(self): + xml_path = ( + pathlib.Path(__file__).resolve().parents[1] + / "roboimi/assets/models/manipulators/DianaMed/table_square.xml" + ) + root = ET.parse(xml_path).getroot() + cameras = {camera.attrib["name"]: camera.attrib for camera in root.findall(".//camera")} + + self.assertNotIn("angle", cameras, "DianaMed scene should stop exposing the old angle camera") + self.assertNotIn("left_side", cameras, "DianaMed scene should no longer expose left_side") + self.assertIn("top", cameras) + self.assertIn("front", cameras) + self.assertEqual(cameras["top"].get("mode"), "targetbody") + self.assertEqual(cameras["top"].get("target"), "table") + + +class AirInsertResetAndStateHelpersTest(unittest.TestCase): + def test_set_socket_peg_task_state_writes_free_joint_qpos(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + setter = getattr(air_insert_env, "set_socket_peg_task_state", None) + self.assertIsNotNone( + setter, + "Expected roboimi.envs.double_air_insert_env.set_socket_peg_task_state", + ) + + socket_qpos = np.zeros(7, dtype=np.float64) + peg_qpos = np.zeros(7, dtype=np.float64) + + class _FakeJoint: + def __init__(self, qpos): + self.qpos = qpos + + class _FakeData: + def joint(self, name): + if name == "blue_socket_joint": + return _FakeJoint(socket_qpos) + if name == "red_peg_joint": + return _FakeJoint(peg_qpos) + raise AssertionError(f"Unexpected joint name: {name}") + + task_state = { + "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float64), + "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64), + "peg_pos": np.array([0.12, 0.91, 0.46], dtype=np.float64), + "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64), + } + + setter(_FakeData(), task_state) + + np.testing.assert_array_equal( + socket_qpos, + np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64), + ) + np.testing.assert_array_equal( + peg_qpos, + np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64), + ) + + def test_get_socket_peg_env_state_returns_stable_14d_vector(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + getter = getattr(air_insert_env, "get_socket_peg_env_state", None) + self.assertIsNotNone( + getter, + "Expected roboimi.envs.double_air_insert_env.get_socket_peg_env_state", + ) + + socket_qpos = np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64) + peg_qpos = np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64) + + class _FakeJoint: + def __init__(self, qpos): + self.qpos = qpos + + class _FakeData: + def joint(self, name): + if name == "blue_socket_joint": + return _FakeJoint(socket_qpos) + if name == "red_peg_joint": + return _FakeJoint(peg_qpos) + raise AssertionError(f"Unexpected joint name: {name}") + + env_state = getter(_FakeData()) + + self.assertEqual(env_state.shape, (14,)) + np.testing.assert_array_equal( + env_state, + np.array( + [-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0, 0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], + dtype=np.float64, + ), + ) + + def test_air_insert_env_does_not_script_attach_or_assist_objects_after_reset(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + env_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None) + self.assertIsNotNone(env_cls) + + source = inspect.getsource(env_cls) + + self.assertNotIn("_update_scripted_grasped_objects", source) + self.assertNotIn("_scripted_", source) + self.assertNotIn("_stabilize_ring_grasp", source) + self.assertNotIn("_ring_grasp_locked", source) + get_reward_source = inspect.getsource(env_cls._get_reward) + self.assertNotIn("ring_block", get_reward_source) + self.assertNotIn("bar_block", get_reward_source) + + def test_socket_peg_xml_defines_active_socket_and_peg_objects(self): + asset_dir = pathlib.Path(__file__).resolve().parents[1] / "roboimi/assets/models/manipulators/DianaMed" + xml_path = asset_dir / "socket_peg_objects.xml" + self.assertTrue(xml_path.exists(), "socket/peg objects should live in socket_peg_objects.xml") + self.assertFalse((asset_dir / "ring_bar_objects.xml").exists(), "old ring_bar_objects.xml should be renamed") + + root = ET.parse(xml_path).getroot() + body_names = {body.attrib.get("name") for body in root.findall(".//body")} + geom_names = {geom.attrib.get("name") for geom in root.findall(".//geom")} + joint_names = {joint.attrib.get("name") for joint in root.findall(".//joint")} + + self.assertIn("socket", body_names) + self.assertIn("peg", body_names) + self.assertNotIn("ring_block", body_names) + self.assertNotIn("bar_block", body_names) + self.assertIn("blue_socket_joint", joint_names) + self.assertIn("red_peg_joint", joint_names) + for geom_name in ("socket-1", "socket-2", "socket-3", "socket-4", "pin", "red_peg"): + self.assertIn(geom_name, geom_names) + + def test_socket_peg_wrapper_includes_socket_peg_objects(self): + xml_path = ( + pathlib.Path(__file__).resolve().parents[1] + / "roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml" + ) + self.assertTrue(xml_path.exists(), "socket/peg wrapper XML should use the new task name") + root = ET.parse(xml_path).getroot() + includes = [include.attrib.get("file") for include in root.findall(".//include")] + self.assertIn("./socket_peg_objects.xml", includes) + self.assertNotIn("./ring_bar_objects.xml", includes) + + +class AirInsertRewardAndSuccessTest(unittest.TestCase): + @staticmethod + def _make_env_state( + socket_pos=(0.0, 0.0, 0.472), + socket_quat=(1.0, 0.0, 0.0, 0.0), + peg_pos=(0.0, 0.0, 0.46), + peg_quat=(1.0, 0.0, 0.0, 0.0), + ): + return np.array([*socket_pos, *socket_quat, *peg_pos, *peg_quat], dtype=np.float64) + + def test_compute_air_insert_reward_counts_left_contact_stage(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None) + self.assertIsNotNone(reward_fn) + + reward = reward_fn( + contact_pairs=[ + ("socket-1", "l_finger_left"), + ("socket-1", "table"), + ("red_peg", "table"), + ], + env_state=self._make_env_state(), + ) + + self.assertEqual(reward, 1) + + def test_compute_air_insert_reward_counts_right_contact_stage(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None) + + reward = reward_fn( + contact_pairs=[ + ("socket-1", "l_finger_left"), + ("red_peg", "l_finger_right"), + ("socket-1", "table"), + ("red_peg", "table"), + ], + env_state=self._make_env_state(), + ) + + self.assertEqual(reward, 2) + + def test_compute_air_insert_reward_counts_lift_stages(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None) + + reward = reward_fn( + contact_pairs=[ + ("socket-1", "l_finger_left"), + ("red_peg", "l_finger_right"), + ], + env_state=self._make_env_state(), + ) + + self.assertEqual(reward, 4) + + def test_compute_air_insert_reward_counts_visual_fingertip_contacts_as_gripper_contacts(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None) + + reward = reward_fn( + contact_pairs=[ + ("socket-3", "r_fingertip_g0_vis_left"), + ("red_peg", "l_fingertip_g0_vis_right"), + ], + env_state=self._make_env_state(), + ) + + self.assertEqual( + reward, + 4, + "visual fingertip geoms are collidable in the Diana XML and should count as gripper-object contacts", + ) + + def test_peg_inserted_into_socket_uses_pin_contact(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + success_fn = getattr(air_insert_env, "peg_inserted_into_socket", None) + self.assertIsNotNone( + success_fn, + "Expected roboimi.envs.double_air_insert_env.peg_inserted_into_socket", + ) + + self.assertTrue(success_fn([("red_peg", "pin")])) + self.assertTrue(success_fn([("pin", "red_peg")])) + self.assertFalse(success_fn([("red_peg", "socket-1")])) + + def test_compute_air_insert_reward_requires_airborne_success_for_final_point(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None) + + reward = reward_fn( + contact_pairs=[ + ("socket-1", "l_finger_left"), + ("red_peg", "l_finger_right"), + ("socket-1", "table"), + ("red_peg", "pin"), + ], + env_state=self._make_env_state(), + ) + + self.assertEqual(reward, 3) + + def test_compute_air_insert_reward_returns_full_score_on_true_airborne_insert(self): + air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env") + reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None) + + reward = reward_fn( + contact_pairs=[ + ("socket-1", "l_finger_left"), + ("red_peg", "l_finger_right"), + ("red_peg", "pin"), + ], + env_state=self._make_env_state(), + ) + + self.assertEqual(reward, 5) + + +class AirInsertPolicyAndSmokeTest(unittest.TestCase): + @staticmethod + def _canonical_task_state(): + return { + "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float32), + "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + "peg_pos": np.array([0.12, 0.90, 0.46], dtype=np.float32), + "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + } + + def test_air_insert_policy_emits_valid_16d_action(self): + policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy") + policy_cls = getattr(policy_module, "TestAirInsertPolicy", None) + self.assertIsNotNone(policy_cls) + + task_state = act_ex_utils.sample_air_insert_socket_peg_state() + policy = policy_cls(inject_noise=False) + action = policy.predict(task_state, 0) + + self.assertEqual(action.shape, (16,)) + np.testing.assert_array_equal(action[-2:], np.array([100, 100])) + + def test_air_insert_policy_inserts_peg_front_view_right_to_left_along_world_x(self): + policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy") + policy_cls = getattr(policy_module, "TestAirInsertPolicy", None) + self.assertIsNotNone(policy_cls) + + task_state = self._canonical_task_state() + policy = policy_cls(inject_noise=False) + policy.generate_trajectory(task_state) + + start_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_START_T) + end_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_END_T) + + self.assertLess( + end_waypoint["xyz"][0], + start_waypoint["xyz"][0] - 0.10, + "front-view right-to-left peg insertion should decrease world x substantially", + ) + self.assertAlmostEqual(float(end_waypoint["xyz"][1]), float(start_waypoint["xyz"][1]), delta=0.02) + expected_insert_end_x = float(task_state["socket_pos"][0] + 0.168) + self.assertAlmostEqual(float(end_waypoint["xyz"][0]), expected_insert_end_x, delta=0.02) + self.assertGreater(float(start_waypoint["xyz"][2]), 0.70) + + def test_air_insert_policy_default_left_grasps_socket_and_right_grasps_peg(self): + policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy") + policy_cls = getattr(policy_module, "TestAirInsertPolicy", None) + self.assertIsNotNone(policy_cls) + + task_state = { + "socket_pos": np.array([-0.18, 0.78, 0.472], dtype=np.float32), + "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + "peg_pos": np.array([0.16, 0.98, 0.46], dtype=np.float32), + "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + } + + policy = policy_cls(inject_noise=False) + policy.generate_trajectory(task_state) + left_close = next(wp for wp in policy.left_trajectory if wp["t"] == 180) + right_close = next(wp for wp in policy.right_trajectory if wp["t"] == 180) + action_z_offset = getattr(policy_cls, "ACTION_OBJECT_Z_OFFSET", 0.11) + expected_socket_pick = task_state["socket_pos"] + np.array([-0.078, 0.0, action_z_offset]) + expected_peg_pick = task_state["peg_pos"] + np.array([0.078, 0.0, action_z_offset + 0.01]) + + np.testing.assert_allclose(left_close["xyz"], expected_socket_pick, atol=1e-6) + np.testing.assert_allclose(right_close["xyz"], expected_peg_pick, atol=1e-6) + self.assertLess(left_close["gripper"], 0, "default policy should close the left gripper on the socket") + self.assertLess(right_close["gripper"], 0, "default policy should close the right gripper on the peg") + + def test_air_insert_policy_socket_hold_tracks_socket_xy_without_sweeping_laterally(self): + policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy") + policy_cls = getattr(policy_module, "TestAirInsertPolicy", None) + self.assertIsNotNone(policy_cls) + + base_state = { + "socket_pos": np.array([-0.20, 0.72, 0.472], dtype=np.float32), + "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + "peg_pos": np.array([0.14, 0.76, 0.46], dtype=np.float32), + "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + } + shifted_state = dict(base_state) + shifted_state["socket_pos"] = np.array([-0.06, 0.99, 0.472], dtype=np.float32) + + base_policy = policy_cls(inject_noise=False) + base_policy.generate_trajectory(base_state) + shifted_policy = policy_cls(inject_noise=False) + shifted_policy.generate_trajectory(shifted_state) + + base_hold = next(wp for wp in base_policy.left_trajectory if wp["t"] == 450) + shifted_hold = next(wp for wp in shifted_policy.left_trajectory if wp["t"] == 450) + np.testing.assert_allclose( + base_hold["xyz"][:2], + base_state["socket_pos"][:2] + np.array([-0.078, 0.0]), + atol=1e-6, + ) + np.testing.assert_allclose( + shifted_hold["xyz"][:2], + shifted_state["socket_pos"][:2] + np.array([-0.078, 0.0]), + atol=1e-6, + ) + + def test_air_insert_policy_predicts_through_full_episode_without_exhausting_waypoints(self): + policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy") + policy_cls = getattr(policy_module, "TestAirInsertPolicy", None) + self.assertIsNotNone(policy_cls) + + task_state = self._canonical_task_state() + policy = policy_cls(inject_noise=False) + + for step in range(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"]): + action = policy.predict(task_state, step) + self.assertEqual(action.shape, (16,)) + + def test_scripted_rollout_entrypoint_selects_socket_peg_sampler_and_policy(self): + rollout_module = importlib.import_module("roboimi.demos.diana_record_sim_episodes") + sampler_fn = getattr(rollout_module, "sample_task_state", None) + policy_factory = getattr(rollout_module, "make_policy", None) + self.assertIsNotNone(sampler_fn) + self.assertIsNotNone(policy_factory) + + task_state = sampler_fn(TASK_NAME) + self.assertEqual(list(task_state.keys()), ["socket_pos", "socket_quat", "peg_pos", "peg_quat"]) + + policy = policy_factory(TASK_NAME, inject_noise=False) + self.assertEqual(policy.__class__.__name__, "TestAirInsertPolicy") + + def test_real_headless_smoke_instantiates_resets_and_steps_new_task_once(self): + policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy") + policy_cls = getattr(policy_module, "TestAirInsertPolicy", None) + self.assertIsNotNone(policy_cls) + + task_state = act_ex_utils.sample_air_insert_socket_peg_state() + env = make_sim_env(TASK_NAME, headless=True) + policy = policy_cls(inject_noise=False) + + try: + env.reset(task_state) + action = policy.predict(task_state, 0) + env.step(action) + self.assertIsNotNone(env.obs) + self.assertIn("qpos", env.obs) + self.assertIn("images", env.obs) + finally: + env.exit_flag = True + cam_thread = getattr(env, "cam_thread", None) + if cam_thread is not None: + cam_thread.join(timeout=1.0) + viewer = getattr(env, "viewer", None) + if viewer is not None: + viewer.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_attnres_resnet2d_backbone.py b/tests/test_attnres_resnet2d_backbone.py new file mode 100644 index 0000000..3bb1829 --- /dev/null +++ b/tests/test_attnres_resnet2d_backbone.py @@ -0,0 +1,26 @@ +import unittest + +import torch + + +class AttnResResNet2DBackboneTest(unittest.TestCase): + def test_backbone_preserves_resnet_like_stage_contract(self): + from roboimi.vla.models.backbones.attnres_resnet2d import AttnResResNetLikeBackbone2D + + backbone = AttnResResNetLikeBackbone2D( + input_channels=3, + stem_dim=16, + stage_dims=(16, 32, 64, 128), + stage_depths=(1, 1, 1, 1), + stage_heads=(2, 4, 4, 8), + stage_kv_heads=(1, 1, 1, 1), + stage_window_sizes=(7, 7, 7, 7), + dropout=0.0, + ) + x = torch.randn(2, 3, 56, 56) + y = backbone(x) + self.assertEqual(y.shape, (2, 128, 2, 2)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_calculate_stats_cli.py b/tests/test_calculate_stats_cli.py new file mode 100644 index 0000000..a298422 --- /dev/null +++ b/tests/test_calculate_stats_cli.py @@ -0,0 +1,88 @@ +import pickle +import tempfile +import unittest +from pathlib import Path + +import h5py +import numpy as np + +from roboimi.vla.scripts import calculate_stats + + +class CalculateStatsCliTest(unittest.TestCase): + def test_default_dataset_dir_is_absolute_and_package_relative(self): + expected = ( + Path(calculate_stats.__file__).resolve().parents[2] + / "demos" + / "dataset" + / "sim_transfer" + ) + + self.assertEqual(Path(calculate_stats.DEFAULT_DATASET_DIR), expected) + self.assertTrue(Path(calculate_stats.DEFAULT_DATASET_DIR).is_absolute()) + + def test_main_writes_dataset_stats_pkl_to_dataset_dir(self): + with tempfile.TemporaryDirectory() as tmpdir: + dataset_dir = Path(tmpdir) + episode_path = dataset_dir / "episode_0.hdf5" + + with h5py.File(episode_path, "w") as root: + root.create_dataset( + "action", + data=np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + ) + observations = root.create_group("observations") + observations.create_dataset( + "qpos", + data=np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32), + ) + + calculate_stats.main(["--dataset_dir", str(dataset_dir)]) + + stats_path = dataset_dir / "dataset_stats.pkl" + self.assertTrue(stats_path.exists()) + + with stats_path.open("rb") as f: + stats = pickle.load(f) + + self.assertEqual( + set(stats), + { + "action_mean", + "action_std", + "action_min", + "action_max", + "qpos_mean", + "qpos_std", + "qpos_min", + "qpos_max", + }, + ) + np.testing.assert_allclose(stats["action_mean"], np.array([2.0, 3.0])) + np.testing.assert_allclose(stats["qpos_mean"], np.array([6.0, 7.0])) + + def test_main_raises_clear_error_for_empty_dataset_dir(self): + with tempfile.TemporaryDirectory() as tmpdir: + dataset_dir = Path(tmpdir) + + with self.assertRaisesRegex( + ValueError, r"No episode_\*\.hdf5 files found" + ) as ctx: + calculate_stats.main(["--dataset_dir", str(dataset_dir)]) + + self.assertIn(str(dataset_dir), str(ctx.exception)) + + def test_main_raises_clear_error_for_missing_dataset_dir(self): + with tempfile.TemporaryDirectory() as tmpdir: + dataset_dir = Path(tmpdir) / "missing" + + with self.assertRaisesRegex( + ValueError, r"No episode_\*\.hdf5 files found" + ) as ctx: + calculate_stats.main(["--dataset_dir", str(dataset_dir)]) + + self.assertIn(str(dataset_dir), str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_eval_vla_execution.py b/tests/test_eval_vla_execution.py new file mode 100644 index 0000000..6a468ac --- /dev/null +++ b/tests/test_eval_vla_execution.py @@ -0,0 +1,28 @@ +import unittest + +from roboimi.vla.eval_utils import execute_policy_action + + +class _FakeEnv: + def __init__(self): + self.calls = [] + + def step(self, action): + self.calls.append(("step", action)) + + def step_jnt(self, action): + self.calls.append(("step_jnt", action)) + + +class EvalVLAExecutionTest(unittest.TestCase): + def test_execute_policy_action_uses_ee_step(self): + env = _FakeEnv() + action = [1, 2, 3] + + execute_policy_action(env, action) + + self.assertEqual(env.calls, [("step", action)]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_eval_vla_headless.py b/tests/test_eval_vla_headless.py new file mode 100644 index 0000000..416c6cf --- /dev/null +++ b/tests/test_eval_vla_headless.py @@ -0,0 +1,391 @@ +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np +import torch +from omegaconf import OmegaConf + +from roboimi.demos.vla_scripts import eval_vla +from roboimi.envs.double_base import DualDianaMed +from roboimi.envs.double_pos_ctrl_env import make_sim_env + + +class _FakeAgent: + def __init__(self): + self.reset_calls = 0 + self.last_observation = None + + def eval(self): + return self + + def to(self, _device): + return self + + def reset(self): + self.reset_calls += 1 + + def select_action(self, observation): + self.last_observation = observation + return torch.zeros(16) + + +class _FakeEnv: + def __init__(self): + self.image_obs_calls = 0 + self.render_calls = 0 + self.reset_calls = [] + + def reset(self, task_state): + self.reset_calls.append(task_state) + + def _get_image_obs(self): + self.image_obs_calls += 1 + return { + "images": { + "front": np.zeros((8, 8, 3), dtype=np.uint8), + } + } + + def _get_qpos_obs(self): + return {"qpos": np.zeros(16, dtype=np.float32)} + + def render(self): + self.render_calls += 1 + raise AssertionError("env.render() should be skipped when eval.headless=true") + + +class _RewardTrackingEnv(_FakeEnv): + def __init__(self, reward_sequences): + super().__init__() + self.reward_sequences = reward_sequences + self.episode_index = -1 + self.step_index = 0 + self.rew = 0.0 + + def reset(self, box_pos): + super().reset(box_pos) + self.episode_index += 1 + self.step_index = 0 + + +class _FakeRenderer: + def __init__(self, env): + self._env = env + self._frames = [ + np.full((4, 4, 3), fill_value=index, dtype=np.uint8) + for index in range(8) + ] + self._index = 0 + + def update_scene(self, _mj_data, camera=None): + self._camera = camera + + def render(self): + frame = self._frames[self._index] + self._index += 1 + if self._index >= len(self._frames): + self._env.exit_flag = True + return frame + + +class EvalVLAHeadlessTest(unittest.TestCase): + def test_prepare_observation_skips_resize_when_image_resize_shape_is_none(self): + obs = { + "images": { + "front": np.arange(8 * 8 * 3, dtype=np.uint8).reshape(8, 8, 3), + }, + "qpos": np.zeros(16, dtype=np.float32), + } + + with mock.patch("cv2.resize", side_effect=AssertionError("resize should be skipped")): + prepared = eval_vla.prepare_observation( + obs, + ["front"], + image_resize_shape=None, + ) + + self.assertEqual(tuple(prepared["images"]["front"].shape), (3, 8, 8)) + self.assertEqual(tuple(prepared["qpos"].shape), (16,)) + + def test_headless_eval_sets_mujoco_gl_to_egl_when_display_missing(self): + cfg = OmegaConf.create({"eval": {"headless": True}}) + with mock.patch.dict(eval_vla.os.environ, {}, clear=True): + eval_vla._configure_headless_mujoco_gl(cfg.eval) + self.assertEqual(eval_vla.os.environ.get("MUJOCO_GL"), "egl") + + def test_headless_eval_preserves_existing_mujoco_gl(self): + cfg = OmegaConf.create({"eval": {"headless": True}}) + with mock.patch.dict(eval_vla.os.environ, {"MUJOCO_GL": "osmesa"}, clear=True): + eval_vla._configure_headless_mujoco_gl(cfg.eval) + self.assertEqual(eval_vla.os.environ.get("MUJOCO_GL"), "osmesa") + + def test_eval_config_exposes_headless_default(self): + eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml")) + + self.assertIn("headless", eval_cfg) + self.assertFalse(eval_cfg.headless) + + def test_make_sim_env_accepts_headless_and_disables_render(self): + fake_env = object() + + with mock.patch( + "roboimi.assets.robots.diana_med.BiDianaMed", + return_value="robot", + ), mock.patch( + "roboimi.envs.double_pos_ctrl_env.DualDianaMed_Pos_Ctrl", + return_value=fake_env, + ) as env_cls: + env = make_sim_env("sim_transfer", headless=True) + + self.assertIs(env, fake_env) + env_cls.assert_called_once_with( + robot="robot", + is_render=False, + control_freq=30, + is_interpolate=True, + cam_view="top", + ) + + def test_headless_sync_camera_capture_populates_images_without_gui_calls(self): + env = DualDianaMed.__new__(DualDianaMed) + env.mj_model = object() + env.mj_data = object() + env.exit_flag = False + env.is_render = False + env.cam = 'top' + env.r_vis = None + env.l_vis = None + env.top = None + env.front = None + env._offscreen_renderer = None + + with mock.patch( + 'roboimi.envs.double_base.mj.Renderer', + side_effect=lambda *args, **kwargs: _FakeRenderer(env), + ) as renderer_cls, mock.patch('roboimi.envs.double_base.cv2.namedWindow') as named_window, mock.patch( + 'roboimi.envs.double_base.cv2.imshow' + ) as imshow, mock.patch('roboimi.envs.double_base.cv2.waitKey') as wait_key: + env._update_camera_images_sync() + + renderer_cls.assert_called_once() + named_window.assert_not_called() + imshow.assert_not_called() + wait_key.assert_not_called() + self.assertIsNotNone(env.r_vis) + self.assertIsNotNone(env.l_vis) + self.assertIsNotNone(env.top) + self.assertIsNotNone(env.front) + + def test_cam_start_skips_background_thread_when_headless(self): + env = DualDianaMed.__new__(DualDianaMed) + env.is_render = False + env.cam_thread = None + + with mock.patch('roboimi.envs.double_base.threading.Thread') as thread_cls: + env.cam_start() + + thread_cls.assert_not_called() + self.assertIsNone(env.cam_thread) + + def test_camera_viewer_headless_updates_images_without_gui_calls(self): + env = DualDianaMed.__new__(DualDianaMed) + env.mj_model = object() + env.mj_data = object() + env.exit_flag = False + env.is_render = False + env.cam = "top" + env.r_vis = None + env.l_vis = None + env.top = None + env.front = None + + with mock.patch( + "roboimi.envs.double_base.mj.Renderer", + side_effect=lambda *args, **kwargs: _FakeRenderer(env), + ), mock.patch("roboimi.envs.double_base.cv2.namedWindow") as named_window, mock.patch( + "roboimi.envs.double_base.cv2.imshow" + ) as imshow, mock.patch("roboimi.envs.double_base.cv2.waitKey") as wait_key: + env.camera_viewer() + + named_window.assert_not_called() + imshow.assert_not_called() + wait_key.assert_not_called() + self.assertIsNotNone(env.r_vis) + self.assertIsNotNone(env.l_vis) + self.assertIsNotNone(env.top) + self.assertIsNotNone(env.front) + + def test_eval_main_headless_skips_render_and_still_executes_policy(self): + fake_env = _FakeEnv() + fake_agent = _FakeAgent() + cfg = OmegaConf.create( + { + "agent": {}, + "eval": { + "ckpt_path": "checkpoints/vla_model_best.pt", + "num_episodes": 1, + "max_timesteps": 1, + "device": "cpu", + "task_name": "sim_transfer", + "camera_names": ["front"], + "use_smoothing": False, + "smooth_alpha": 0.3, + "verbose_action": False, + "headless": True, + }, + } + ) + + with mock.patch.object( + eval_vla, + "load_checkpoint", + return_value=(fake_agent, None), + ), mock.patch.object( + eval_vla, + "make_sim_env", + return_value=fake_env, + ) as make_env, mock.patch.object( + eval_vla, + "sample_transfer_pose", + return_value=np.array([0.1, 0.2, 0.3]), + ), mock.patch.object( + eval_vla, + "execute_policy_action", + ) as execute_policy_action, mock.patch.object( + eval_vla, + "tqdm", + side_effect=lambda iterable, **kwargs: iterable, + ): + eval_vla.main.__wrapped__(cfg) + + make_env.assert_called_once_with("sim_transfer", headless=True) + execute_policy_action.assert_called_once() + self.assertEqual(fake_env.image_obs_calls, 1) + self.assertEqual(fake_env.render_calls, 0) + self.assertIsNotNone(fake_agent.last_observation) + self.assertIn("front", fake_agent.last_observation["images"]) + + def test_run_eval_returns_average_reward_summary(self): + reward_sequences = [ + [1.0, 2.0], + [0.5, 4.0], + ] + fake_env = _RewardTrackingEnv(reward_sequences) + fake_agent = _FakeAgent() + cfg = OmegaConf.create( + { + "agent": {}, + "eval": { + "ckpt_path": "checkpoints/vla_model_best.pt", + "num_episodes": 2, + "max_timesteps": 2, + "device": "cpu", + "task_name": "sim_transfer", + "camera_names": ["front"], + "use_smoothing": False, + "smooth_alpha": 0.3, + "verbose_action": False, + "headless": True, + }, + } + ) + + def fake_execute_policy_action(env, action): + del action + env.rew = env.reward_sequences[env.episode_index][env.step_index] + env.step_index += 1 + + with mock.patch.object( + eval_vla, + "load_checkpoint", + return_value=(fake_agent, None), + ), mock.patch.object( + eval_vla, + "make_sim_env", + return_value=fake_env, + ), mock.patch.object( + eval_vla, + "sample_transfer_pose", + return_value=np.array([0.1, 0.2, 0.3]), + ), mock.patch.object( + eval_vla, + "execute_policy_action", + side_effect=fake_execute_policy_action, + ), mock.patch.object( + eval_vla, + "tqdm", + side_effect=lambda iterable, **kwargs: iterable, + ): + summary = eval_vla._run_eval(cfg) + + self.assertEqual(summary["episode_rewards"], [3.0, 4.5]) + self.assertAlmostEqual(summary["avg_reward"], 3.75) + self.assertEqual(summary["num_episodes"], 2) + + def test_run_eval_uses_air_insert_sampler_for_socket_peg_task(self): + self.assertTrue( + hasattr(eval_vla, "sample_air_insert_socket_peg_state"), + "Expected eval_vla to expose the new socket/peg reset sampler", + ) + + fake_env = _FakeEnv() + fake_agent = _FakeAgent() + sampled_task_state = { + "socket_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32), + "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + "peg_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32), + "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32), + } + cfg = OmegaConf.create( + { + "agent": {}, + "eval": { + "ckpt_path": "checkpoints/vla_model_best.pt", + "num_episodes": 1, + "max_timesteps": 1, + "device": "cpu", + "task_name": "sim_air_insert_socket_peg", + "camera_names": ["front"], + "use_smoothing": False, + "smooth_alpha": 0.3, + "verbose_action": False, + "headless": True, + }, + } + ) + + with mock.patch.object( + eval_vla, + "load_checkpoint", + return_value=(fake_agent, None), + ), mock.patch.object( + eval_vla, + "make_sim_env", + return_value=fake_env, + ) as make_env, mock.patch.object( + eval_vla, + "sample_air_insert_socket_peg_state", + return_value=sampled_task_state, + ) as socket_peg_sampler, mock.patch.object( + eval_vla, + "sample_transfer_pose", + side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_socket_peg"), + ), mock.patch.object( + eval_vla, + "execute_policy_action", + ) as execute_policy_action, mock.patch.object( + eval_vla, + "tqdm", + side_effect=lambda iterable, **kwargs: iterable, + ): + eval_vla._run_eval(cfg) + + make_env.assert_called_once_with("sim_air_insert_socket_peg", headless=True) + socket_peg_sampler.assert_called_once_with() + execute_policy_action.assert_called_once() + self.assertEqual(fake_env.reset_calls, [sampled_task_state]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_eval_vla_headless_import.py b/tests/test_eval_vla_headless_import.py new file mode 100644 index 0000000..e9d4763 --- /dev/null +++ b/tests/test_eval_vla_headless_import.py @@ -0,0 +1,26 @@ +import json +import os +import subprocess +import sys + + +def test_eval_vla_import_does_not_import_mujoco_early_when_headless_backend_not_set(): + env = os.environ.copy() + env.pop('MUJOCO_GL', None) + proc = subprocess.run( + [ + sys.executable, + '-c', + ( + 'import json, sys; ' + 'from roboimi.demos.vla_scripts import eval_vla; ' + 'print(json.dumps({"mujoco_in_sys_modules": "mujoco" in sys.modules}))' + ), + ], + capture_output=True, + text=True, + env=env, + check=True, + ) + payload = json.loads(proc.stdout.strip()) + assert payload['mujoco_in_sys_modules'] is False diff --git a/tests/test_eval_vla_rollout_artifacts.py b/tests/test_eval_vla_rollout_artifacts.py new file mode 100644 index 0000000..7cb316a --- /dev/null +++ b/tests/test_eval_vla_rollout_artifacts.py @@ -0,0 +1,344 @@ +import json +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np +import torch +from omegaconf import OmegaConf + +from roboimi.demos.vla_scripts import eval_vla + + +class _FakeAgent: + def __init__(self, actions): + self._actions = [torch.tensor(action, dtype=torch.float32) for action in actions] + self.reset_calls = 0 + + def eval(self): + return self + + def to(self, _device): + return self + + def reset(self): + self.reset_calls += 1 + + def select_action(self, observation): + del observation + return self._actions.pop(0) + + +class _FakeEnv: + def __init__(self): + self.step_count = 0 + self.rew = 0.0 + self.render_calls = 0 + self.reset_calls = [] + + def reset(self, box_pos): + self.reset_calls.append(np.array(box_pos, copy=True)) + self.step_count = 0 + self.rew = 0.0 + + def _get_image_obs(self): + frame_value = self.step_count + front = np.full((6, 8, 3), fill_value=frame_value, dtype=np.uint8) + top = np.full((6, 8, 3), fill_value=frame_value + 20, dtype=np.uint8) + return {"images": {"front": front, "top": top}} + + def _get_qpos_obs(self): + return {"qpos": np.arange(16, dtype=np.float32)} + + def step(self, action): + del action + self.step_count += 1 + self.rew = float(self.step_count) + + def render(self): + self.render_calls += 1 + + def getBodyPos(self, name): + base = float(self.step_count) + if name == 'eef_left': + return np.array([base, base + 0.1, base + 0.2], dtype=np.float32) + if name == 'eef_right': + return np.array([base + 1.0, base + 1.1, base + 1.2], dtype=np.float32) + raise KeyError(name) + + def getBodyQuat(self, name): + base = float(self.step_count) + if name == 'eef_left': + return np.array([1.0, base, 0.0, 0.0], dtype=np.float32) + if name == 'eef_right': + return np.array([1.0, 0.0, base, 0.0], dtype=np.float32) + raise KeyError(name) + + +class _FakeVideoWriter: + def __init__(self, output_path): + self.output_path = Path(output_path) + self.output_path.parent.mkdir(parents=True, exist_ok=True) + self.output_path.write_bytes(b'') + self.frames = [] + self.released = False + + def isOpened(self): + return True + + def write(self, frame): + self.frames.append(np.array(frame, copy=True)) + + def release(self): + self.released = True + self.output_path.write_bytes(b'fake-mp4') + + +class EvalVLARolloutArtifactsTest(unittest.TestCase): + def test_eval_config_exposes_rollout_artifact_defaults(self): + eval_cfg = OmegaConf.load(Path('roboimi/vla/conf/eval/eval.yaml')) + + self.assertIn('artifact_dir', eval_cfg) + self.assertFalse(eval_cfg.save_summary_json) + self.assertFalse(eval_cfg.save_trajectory_npz) + self.assertFalse(eval_cfg.save_trajectory_image) + self.assertFalse(eval_cfg.record_video) + self.assertIsNone(eval_cfg.artifact_dir) + self.assertIsNone(eval_cfg.trajectory_image_camera_name) + self.assertIsNone(eval_cfg.video_camera_name) + self.assertEqual(eval_cfg.video_fps, 30) + + def test_run_eval_exports_npz_summary_and_video_artifacts(self): + actions = [ + np.arange(16, dtype=np.float32), + np.arange(16, dtype=np.float32) + 10.0, + ] + fake_agent = _FakeAgent(actions) + fake_env = _FakeEnv() + + with tempfile.TemporaryDirectory() as tmpdir: + cfg = OmegaConf.create( + { + 'agent': {}, + 'eval': { + 'ckpt_path': 'checkpoints/vla_model_best.pt', + 'num_episodes': 1, + 'max_timesteps': 2, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front', 'top'], + 'use_smoothing': True, + 'smooth_alpha': 0.5, + 'verbose_action': False, + 'headless': True, + 'artifact_dir': tmpdir, + 'save_summary_json': True, + 'save_trajectory_npz': True, + 'save_trajectory_image': True, + 'trajectory_image_camera_name': 'front', + 'record_video': True, + 'video_camera_name': 'front', + 'video_fps': 12, + }, + } + ) + + writer_holder = {} + + def fake_open_video_writer(output_path, frame_size, fps): + self.assertEqual(frame_size, (8, 6)) + self.assertEqual(fps, 12) + writer = _FakeVideoWriter(output_path) + writer_holder['writer'] = writer + return writer + + with mock.patch.object( + eval_vla, + 'load_checkpoint', + return_value=(fake_agent, None), + ), mock.patch.object( + eval_vla, + 'make_sim_env', + return_value=fake_env, + ), mock.patch.object( + eval_vla, + 'sample_transfer_pose', + return_value=np.array([0.1, 0.2, 0.3], dtype=np.float32), + ), mock.patch.object( + eval_vla, + 'tqdm', + side_effect=lambda iterable, **kwargs: iterable, + ), mock.patch.object( + eval_vla, + '_open_video_writer', + side_effect=fake_open_video_writer, + ): + summary = eval_vla._run_eval(cfg) + + artifacts = summary['artifacts'] + trajectory_path = Path(artifacts['trajectory_npz']) + summary_path = Path(artifacts['summary_json']) + video_path = Path(artifacts['video_mp4']) + trajectory_image_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image']) + + self.assertEqual(Path(artifacts['output_dir']), Path(tmpdir)) + self.assertEqual(artifacts['video_camera_name'], 'front') + self.assertTrue(trajectory_path.exists()) + self.assertTrue(summary_path.exists()) + self.assertTrue(video_path.exists()) + self.assertTrue(trajectory_image_path.exists()) + + rollout_npz = np.load(trajectory_path) + np.testing.assert_array_equal(rollout_npz['episode_index'], np.array([0, 0])) + np.testing.assert_array_equal(rollout_npz['timestep'], np.array([0, 1])) + np.testing.assert_array_equal(rollout_npz['reward'], np.array([1.0, 2.0], dtype=np.float32)) + np.testing.assert_array_equal(rollout_npz['raw_predicted_ee_action'][0], actions[0]) + np.testing.assert_array_equal(rollout_npz['raw_predicted_ee_action'][1], actions[1]) + np.testing.assert_array_equal(rollout_npz['executed_ee_action'][0], actions[0]) + np.testing.assert_array_equal( + rollout_npz['executed_ee_action'][1], + (actions[0] + actions[1]) / 2.0, + ) + np.testing.assert_array_equal( + rollout_npz['left_ee_pos'], + np.array([[1.0, 1.1, 1.2], [2.0, 2.1, 2.2]], dtype=np.float32), + ) + np.testing.assert_array_equal( + rollout_npz['right_ee_pos'], + np.array([[2.0, 2.1, 2.2], [3.0, 3.1, 3.2]], dtype=np.float32), + ) + self.assertEqual(rollout_npz['obs_read_time_ms'].shape, (2,)) + self.assertEqual(rollout_npz['preprocess_time_ms'].shape, (2,)) + self.assertEqual(rollout_npz['inference_time_ms'].shape, (2,)) + self.assertEqual(rollout_npz['env_step_time_ms'].shape, (2,)) + self.assertEqual(rollout_npz['total_time_ms'].shape, (2,)) + + writer = writer_holder['writer'] + self.assertTrue(writer.released) + self.assertEqual(len(writer.frames), 2) + np.testing.assert_array_equal(writer.frames[0], np.zeros((6, 8, 3), dtype=np.uint8)) + np.testing.assert_array_equal(writer.frames[1], np.full((6, 8, 3), 1, dtype=np.uint8)) + + with summary_path.open('r', encoding='utf-8') as fh: + saved_summary = json.load(fh) + self.assertEqual(saved_summary['artifacts']['trajectory_npz'], str(trajectory_path)) + self.assertEqual(saved_summary['artifacts']['video_mp4'], str(video_path)) + self.assertEqual( + saved_summary['episodes'][0]['artifact_paths']['trajectory_image'], + str(trajectory_image_path), + ) + self.assertEqual(saved_summary['episode_rewards'], [3.0]) + self.assertAlmostEqual(summary['avg_reward'], 3.0) + self.assertIn('avg_obs_read_time_ms', summary) + self.assertIn('avg_env_step_time_ms', summary) + + def test_run_eval_exports_front_trajectory_images_without_video_dependency(self): + actions = [ + np.arange(16, dtype=np.float32), + np.arange(16, dtype=np.float32) + 10.0, + np.arange(16, dtype=np.float32) + 100.0, + np.arange(16, dtype=np.float32) + 110.0, + ] + fake_agent = _FakeAgent(actions) + fake_env = _FakeEnv() + + with tempfile.TemporaryDirectory() as tmpdir: + cfg = OmegaConf.create( + { + 'agent': {}, + 'eval': { + 'ckpt_path': 'checkpoints/vla_model_best.pt', + 'num_episodes': 2, + 'max_timesteps': 2, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['top', 'front'], + 'use_smoothing': True, + 'smooth_alpha': 0.5, + 'verbose_action': False, + 'headless': True, + 'artifact_dir': tmpdir, + 'save_trajectory_image': True, + 'record_video': False, + }, + } + ) + + trajectory_image_calls = [] + + def fake_save_rollout_trajectory_image( + env, + output_path, + raw_actions, + camera_name, + *, + line_radius=0.004, + max_markers=1500, + ): + del env, line_radius, max_markers + trajectory_image_calls.append( + { + 'output_path': output_path, + 'camera_name': camera_name, + 'raw_actions': [np.array(action, copy=True) for action in raw_actions], + } + ) + if output_path is None: + return None + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_bytes(b'fake-png') + return str(output_path) + + with mock.patch.object( + eval_vla, + 'load_checkpoint', + return_value=(fake_agent, None), + ), mock.patch.object( + eval_vla, + 'make_sim_env', + return_value=fake_env, + ), mock.patch.object( + eval_vla, + 'sample_transfer_pose', + return_value=np.array([0.1, 0.2, 0.3], dtype=np.float32), + ), mock.patch.object( + eval_vla, + 'tqdm', + side_effect=lambda iterable, **kwargs: iterable, + ), mock.patch.object( + eval_vla, + '_save_rollout_trajectory_image', + side_effect=fake_save_rollout_trajectory_image, + ) as save_trajectory_image_mock, mock.patch.object( + eval_vla, + '_open_video_writer', + ) as open_video_writer_mock: + summary = eval_vla._run_eval(cfg) + + self.assertEqual(save_trajectory_image_mock.call_count, 2) + open_video_writer_mock.assert_not_called() + self.assertIsNone(summary['artifacts']['video_mp4']) + self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front') + self.assertEqual( + [call['camera_name'] for call in trajectory_image_calls], + ['front', 'front'], + ) + + first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image']) + second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image']) + self.assertTrue(first_episode_path.exists()) + self.assertTrue(second_episode_path.exists()) + self.assertNotEqual(first_episode_path, second_episode_path) + self.assertEqual(first_episode_path.parent, Path(tmpdir)) + self.assertEqual(second_episode_path.parent, Path(tmpdir)) + + np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0]) + np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1]) + np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2]) + np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3]) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_imf_transformer1d_external_alignment.py b/tests/test_imf_transformer1d_external_alignment.py new file mode 100644 index 0000000..c3dff34 --- /dev/null +++ b/tests/test_imf_transformer1d_external_alignment.py @@ -0,0 +1,196 @@ +import contextlib +import importlib +import inspect +import subprocess +import sys +import types +import unittest +from pathlib import Path + +import torch + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +_EXTERNAL_COMMIT = '185ed659' +_LOCAL_MODULE_NAME = 'roboimi.vla.models.heads.imf_transformer1d' +_MISSING = object() + + +def _find_external_checkout_root() -> Path | None: + for ancestor in (_REPO_ROOT, *_REPO_ROOT.parents): + candidate = ancestor / 'diffusion_policy' + if (candidate / '.git').exists(): + return candidate + return None + + +_EXTERNAL_CHECKOUT_ROOT = _find_external_checkout_root() +_EXTERNAL_MODULE_PATHS = { + 'diffusion_policy.model.common.module_attr_mixin': 'diffusion_policy/model/common/module_attr_mixin.py', + 'diffusion_policy.model.diffusion.positional_embedding': 'diffusion_policy/model/diffusion/positional_embedding.py', + 'diffusion_policy.model.diffusion.attnres_transformer_components': 'diffusion_policy/model/diffusion/attnres_transformer_components.py', + 'diffusion_policy.model.diffusion.imf_transformer_for_diffusion': 'diffusion_policy/model/diffusion/imf_transformer_for_diffusion.py', +} + + +@contextlib.contextmanager +def _temporary_registered_modules(): + previous_modules = {} + + def remember(name: str) -> None: + if name not in previous_modules: + previous_modules[name] = sys.modules.get(name, _MISSING) + + def ensure_package(name: str) -> None: + if not name or name in sys.modules: + return + remember(name) + package = types.ModuleType(name) + package.__path__ = [] + sys.modules[name] = package + + def load(name: str, source: str, origin: str): + package_parts = name.split('.')[:-1] + for idx in range(1, len(package_parts) + 1): + ensure_package('.'.join(package_parts[:idx])) + + remember(name) + module = types.ModuleType(name) + module.__file__ = origin + module.__package__ = name.rpartition('.')[0] + sys.modules[name] = module + exec(compile(source, origin, 'exec'), module.__dict__) + return module + + try: + yield load + finally: + for name, previous in reversed(list(previous_modules.items())): + if previous is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + +def _git_show(repo_root: Path, commit: str, relative_path: str) -> str: + result = subprocess.run( + ['git', '-C', str(repo_root), 'show', f'{commit}:{relative_path}'], + check=True, + capture_output=True, + text=True, + ) + return result.stdout + + +@contextlib.contextmanager +def _load_external_module_or_skip(test_case: unittest.TestCase): + if _EXTERNAL_CHECKOUT_ROOT is None: + test_case.skipTest('external diffusion_policy checkout unavailable') + + try: + sources = { + name: _git_show(_EXTERNAL_CHECKOUT_ROOT, _EXTERNAL_COMMIT, relative_path) + for name, relative_path in _EXTERNAL_MODULE_PATHS.items() + } + except subprocess.CalledProcessError as exc: + test_case.skipTest( + f'external diffusion_policy commit {_EXTERNAL_COMMIT} is unavailable: {exc.stderr.strip() or exc}' + ) + + with _temporary_registered_modules() as load_external: + for name, relative_path in _EXTERNAL_MODULE_PATHS.items(): + load_external( + name, + sources[name], + origin=f'{_EXTERNAL_CHECKOUT_ROOT}:{_EXTERNAL_COMMIT}:{relative_path}', + ) + yield sys.modules['diffusion_policy.model.diffusion.imf_transformer_for_diffusion'] + + +def _load_local_module(): + importlib.invalidate_caches() + sys.modules.pop(_LOCAL_MODULE_NAME, None) + return importlib.import_module(_LOCAL_MODULE_NAME) + + +class IMFTransformer1DExternalAlignmentTest(unittest.TestCase): + def _optim_group_names(self, model, groups): + names_by_param = {id(param): name for name, param in model.named_parameters()} + return [ + {names_by_param[id(param)] for param in group['params']} + for group in groups + ] + + def test_local_defaults_preserve_supported_attnres_config(self): + local_module = _load_local_module() + ctor = inspect.signature(local_module.IMFTransformer1D.__init__).parameters + + self.assertEqual(ctor['backbone_type'].default, 'attnres_full') + self.assertEqual(ctor['n_head'].default, 1) + self.assertEqual(ctor['n_kv_head'].default, 1) + self.assertEqual(ctor['n_cond_layers'].default, 0) + self.assertTrue(ctor['time_as_cond'].default) + self.assertFalse(ctor['causal_attn'].default) + + def test_attnres_full_state_dict_forward_and_optim_groups_match_external(self): + local_module = _load_local_module() + with _load_external_module_or_skip(self) as external_module: + config = dict( + input_dim=4, + output_dim=4, + horizon=6, + n_obs_steps=3, + cond_dim=5, + n_layer=2, + n_head=1, + n_emb=16, + p_drop_emb=0.0, + p_drop_attn=0.0, + causal_attn=False, + time_as_cond=True, + n_cond_layers=0, + backbone_type='attnres_full', + n_kv_head=1, + ) + + torch.manual_seed(7) + external_model = external_module.IMFTransformerForDiffusion(**config) + local_model = local_module.IMFTransformer1D(**config) + external_model.eval() + local_model.eval() + + external_state_dict = external_model.state_dict() + self.assertEqual(set(local_model.state_dict().keys()), set(external_state_dict.keys())) + local_model.load_state_dict(external_state_dict, strict=True) + + batch_size = 2 + sample = torch.randn(batch_size, config['horizon'], config['input_dim']) + r = torch.tensor([0.1, 0.4], dtype=torch.float32) + t = torch.tensor([0.7, 0.9], dtype=torch.float32) + cond = torch.randn(batch_size, config['n_obs_steps'], config['cond_dim']) + + with torch.no_grad(): + external_out = external_model(sample=sample, r=r, t=t, cond=cond) + local_out = local_model(sample=sample, r=r, t=t, cond=cond) + + self.assertEqual(local_out.shape, (batch_size, config['horizon'], config['output_dim'])) + self.assertEqual(local_out.shape, external_out.shape) + self.assertTrue(torch.allclose(local_out, external_out, atol=1e-6, rtol=1e-5)) + + weight_decay = 0.123 + external_groups = external_model.get_optim_groups(weight_decay=weight_decay) + local_groups = local_model.get_optim_groups(weight_decay=weight_decay) + + self.assertEqual(len(local_groups), len(external_groups)) + self.assertEqual([group['weight_decay'] for group in local_groups], [weight_decay, 0.0]) + self.assertEqual( + self._optim_group_names(local_model, local_groups), + self._optim_group_names(external_model, external_groups), + ) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_imf_vla_agent.py b/tests/test_imf_vla_agent.py new file mode 100644 index 0000000..dfccdff --- /dev/null +++ b/tests/test_imf_vla_agent.py @@ -0,0 +1,889 @@ +import contextlib +import importlib +import importlib.machinery +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +import torch +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from hydra.utils import instantiate +from omegaconf import OmegaConf +from torch import nn + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve()) +_MISSING = object() +_CAMERA_NAMES = ('r_vis', 'top', 'front') + + +class _FakeScheduler: + def __init__(self, num_train_timesteps=100, **kwargs): + self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps) + self.timesteps = [] + + def add_noise(self, sample, noise, timestep): + return sample + noise + + def set_timesteps(self, num_inference_steps): + self.timesteps = list(range(num_inference_steps - 1, -1, -1)) + + def step(self, noise_pred, timestep, sample): + return types.SimpleNamespace(prev_sample=sample) + + +class _IdentityCrop: + def __init__(self, size): + self.size = size + + def __call__(self, x): + return x + + +class _FakeResNet(nn.Module): + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu1 = nn.ReLU() + self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1, stride=2) + self.relu2 = nn.ReLU() + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(16, 16) + + def forward(self, x): + x = self.relu1(self.conv1(x)) + x = self.relu2(self.conv2(x)) + x = self.avgpool(x) + x = torch.flatten(x, start_dim=1) + return self.fc(x) + + +class _FakeRearrange(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, x): + return x + + +class _FakeViTConfig: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + +class _FakeViTModel(nn.Module): + def __init__(self, config, add_pooling_layer=False): + super().__init__() + del add_pooling_layer + self.config = config + hidden_size = int(getattr(config, 'hidden_size', 192)) + self.proj = nn.Linear(hidden_size, hidden_size) + + def forward(self, pixel_values=None, interpolate_pos_encoding=False, **kwargs): + del interpolate_pos_encoding, kwargs + batch_size = pixel_values.shape[0] + hidden_size = int(getattr(self.config, 'hidden_size', 192)) + seq_len = 2 + last_hidden_state = torch.zeros(batch_size, seq_len, hidden_size, dtype=pixel_values.dtype, device=pixel_values.device) + return types.SimpleNamespace(last_hidden_state=last_hidden_state) + + +class _FakeSiglipVisionOutput: + def __init__(self, pooler_output): + self.pooler_output = pooler_output + + +class _FakeSiglipVisionConfig: + def __init__(self, hidden_size=768, image_size=256): + self.hidden_size = hidden_size + self.image_size = image_size + + +class _FakeSiglipVisionModel(nn.Module): + load_calls = [] + + def __init__(self, hidden_size=768): + super().__init__() + self.config = _FakeSiglipVisionConfig(hidden_size=hidden_size) + self.scale = nn.Parameter(torch.tensor(1.0)) + self.forward_calls = [] + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): + model = cls() + cls.load_calls.append({ + 'pretrained_model_name_or_path': pretrained_model_name_or_path, + 'args': args, + 'kwargs': kwargs, + }) + return model + + def forward(self, pixel_values=None, **kwargs): + self.forward_calls.append({ + 'pixel_values': pixel_values.detach().clone(), + 'kwargs': dict(kwargs), + }) + pooled = pixel_values.mean(dim=(2, 3), keepdim=False) * self.scale + return _FakeSiglipVisionOutput(pooler_output=pooled) + + +class _StubIMFHead(nn.Module): + def __init__( + self, + input_dim, + output_dim, + horizon, + n_obs_steps, + cond_dim, + **kwargs, + ): + super().__init__() + self.constructor_kwargs = { + 'input_dim': input_dim, + 'output_dim': output_dim, + 'horizon': horizon, + 'n_obs_steps': n_obs_steps, + 'cond_dim': cond_dim, + **kwargs, + } + self.proj = nn.Linear(input_dim, output_dim) + self.cond_obs_emb = nn.Linear(cond_dim, max(cond_dim, 1)) + + def forward(self, sample, r, t, cond=None): + return torch.zeros_like(sample) + + def get_optim_groups(self, weight_decay): + return [ + {'params': [self.proj.weight], 'weight_decay': weight_decay}, + {'params': [self.proj.bias, self.cond_obs_emb.weight, self.cond_obs_emb.bias], 'weight_decay': 0.0}, + ] + + +@contextlib.contextmanager +def _stub_optional_modules(include_imf_head=False): + previous_modules = {} + + def remember_and_remove(name): + if name not in previous_modules: + previous_modules[name] = sys.modules.get(name, _MISSING) + sys.modules.pop(name, None) + + def inject(name, module): + if name not in previous_modules: + previous_modules[name] = sys.modules.get(name, _MISSING) + sys.modules[name] = module + + diffusers_module = types.ModuleType('diffusers') + schedulers_module = types.ModuleType('diffusers.schedulers') + ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm') + ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim') + ddpm_module.DDPMScheduler = _FakeScheduler + ddim_module.DDIMScheduler = _FakeScheduler + diffusers_module.DDPMScheduler = _FakeScheduler + diffusers_module.DDIMScheduler = _FakeScheduler + diffusers_module.schedulers = schedulers_module + schedulers_module.scheduling_ddpm = ddpm_module + schedulers_module.scheduling_ddim = ddim_module + + torchvision_module = types.ModuleType('torchvision') + models_module = types.ModuleType('torchvision.models') + transforms_module = types.ModuleType('torchvision.transforms') + torchvision_module.__spec__ = importlib.machinery.ModuleSpec('torchvision', loader=None) + models_module.__spec__ = importlib.machinery.ModuleSpec('torchvision.models', loader=None) + transforms_module.__spec__ = importlib.machinery.ModuleSpec('torchvision.transforms', loader=None) + models_module.resnet18 = lambda weights=None: _FakeResNet() + transforms_module.CenterCrop = _IdentityCrop + transforms_module.RandomCrop = _IdentityCrop + torchvision_module.models = models_module + torchvision_module.transforms = transforms_module + + einops_module = types.ModuleType('einops') + einops_module.rearrange = lambda x, *args, **kwargs: x + einops_layers_module = types.ModuleType('einops.layers') + einops_layers_torch_module = types.ModuleType('einops.layers.torch') + einops_layers_torch_module.Rearrange = _FakeRearrange + einops_module.layers = einops_layers_module + einops_layers_module.torch = einops_layers_torch_module + + transformers_module = types.ModuleType('transformers') + transformers_module.__spec__ = importlib.machinery.ModuleSpec('transformers', loader=None) + transformers_module.ViTConfig = _FakeViTConfig + transformers_module.ViTModel = _FakeViTModel + transformers_module.SiglipVisionModel = _FakeSiglipVisionModel + + try: + remember_and_remove('roboimi.vla.models.backbones.siglip2_diffusion_backbone') + inject('diffusers', diffusers_module) + inject('diffusers.schedulers', schedulers_module) + inject('diffusers.schedulers.scheduling_ddpm', ddpm_module) + inject('diffusers.schedulers.scheduling_ddim', ddim_module) + inject('torchvision', torchvision_module) + inject('torchvision.models', models_module) + inject('torchvision.transforms', transforms_module) + inject('einops', einops_module) + inject('einops.layers', einops_layers_module) + inject('einops.layers.torch', einops_layers_torch_module) + inject('transformers', transformers_module) + + if include_imf_head: + import roboimi.vla.models.heads as heads_package + + imf_head_module = types.ModuleType('roboimi.vla.models.heads.imf_transformer1d') + imf_head_module.IMFTransformer1D = _StubIMFHead + inject('roboimi.vla.models.heads.imf_transformer1d', imf_head_module) + setattr(heads_package, 'imf_transformer1d', imf_head_module) + + yield + finally: + for name, previous in reversed(list(previous_modules.items())): + if previous is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + +def _compose_cfg(overrides=None): + if not OmegaConf.has_resolver('len'): + OmegaConf.register_new_resolver('len', lambda x: len(x)) + + GlobalHydra.instance().clear() + with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR): + return compose(config_name='config', overrides=list(overrides or [])) + + +def _load_imf_agent_class(): + with _stub_optional_modules(): + sys.modules.pop('roboimi.vla.agent_imf', None) + module = importlib.import_module('roboimi.vla.agent_imf') + return module.IMFVLAAgent, module + + +class _StubVisionBackbone(nn.Module): + output_dim = 1 + + def __init__(self, camera_names=_CAMERA_NAMES): + super().__init__() + self.camera_names = tuple(camera_names) + self.num_cameras = len(self.camera_names) + + def forward(self, images): + per_camera_features = [] + for camera_name in self.camera_names: + image_batch = images[camera_name] + per_camera_features.append(image_batch.mean(dim=(2, 3, 4), keepdim=False).unsqueeze(-1)) + return torch.cat(per_camera_features, dim=-1) + + +class _StubJointVisionBackbone(nn.Module): + joint_output_dim = 5 + output_dim = 5 + + def __init__(self, camera_names=_CAMERA_NAMES): + super().__init__() + self.camera_names = tuple(camera_names) + self.num_cameras = len(self.camera_names) + + def forward(self, images): + batch_size, obs_horizon = next(iter(images.values())).shape[:2] + features = [] + for camera_name in ('front', 'top', 'r_vis'): + image_batch = images[camera_name] + features.append(image_batch.mean(dim=(2, 3, 4), keepdim=False).unsqueeze(-1)) + joint_features = torch.cat(features, dim=-1) + front_top_sum = joint_features[..., :2].sum(dim=-1, keepdim=True) + r_vis_minus_front = (joint_features[..., 2:] - joint_features[..., :1]) + time_marker = torch.arange(obs_horizon, dtype=joint_features.dtype).view(1, obs_horizon, 1) + time_marker = time_marker.expand(batch_size, -1, -1) + return torch.cat([joint_features, front_top_sum, r_vis_minus_front + time_marker], dim=-1) + + +class _StubMultiTokenVisionBackbone(nn.Module): + output_dim = 2 + tokens_per_step = 3 + + def __init__(self, camera_names=_CAMERA_NAMES): + super().__init__() + self.camera_names = tuple(camera_names) + self.num_cameras = len(self.camera_names) + + def forward(self, images): + batch_size, obs_horizon = next(iter(images.values())).shape[:2] + features = [] + time_marker = torch.arange(obs_horizon, dtype=torch.float32).view(1, obs_horizon, 1).expand(batch_size, -1, -1) + for camera_name in self.camera_names: + image_batch = images[camera_name] + camera_marker = image_batch.mean(dim=(2, 3, 4), keepdim=False).unsqueeze(-1) + features.append(torch.cat([camera_marker, camera_marker + time_marker], dim=-1)) + return torch.stack(features, dim=2) + + +class _StubMultiTokenVisionBackbone(nn.Module): + output_dim = 2 + tokens_per_step = 3 + + def __init__(self, camera_names=_CAMERA_NAMES): + super().__init__() + self.camera_names = tuple(camera_names) + self.num_cameras = len(self.camera_names) + + def forward(self, images): + per_camera = [] + for camera_name in self.camera_names: + image_batch = images[camera_name] + base = image_batch.mean(dim=(2, 3, 4), keepdim=False) + per_camera.append(torch.stack([base, base + 0.5], dim=-1)) + return torch.stack(per_camera, dim=2) + + +class _RecordingLinearIMFHead(nn.Module): + def __init__(self): + super().__init__() + self.scale = nn.Parameter(torch.tensor(0.5)) + self.calls = [] + + @staticmethod + def _broadcast_batch_time(value, reference): + while value.ndim < reference.ndim: + value = value.unsqueeze(-1) + return value + + def forward(self, sample, r, t, cond=None): + record = { + 'sample': sample.detach().clone(), + 'r': r.detach().clone(), + 't': t.detach().clone(), + 'cond': None if cond is None else cond.detach().clone(), + } + self.calls.append(record) + cond_term = 0.0 + if cond is not None: + cond_term = cond.mean(dim=(1, 2), keepdim=True) + r_b = self._broadcast_batch_time(r, sample) + t_b = self._broadcast_batch_time(t, sample) + return self.scale * sample + r_b + 2.0 * t_b + cond_term + + +class _ForbiddenScheduler: + def set_timesteps(self, *args, **kwargs): # pragma: no cover - only runs on regression + raise AssertionError('IMF inference should not use DDIM scheduler set_timesteps') + + def step(self, *args, **kwargs): # pragma: no cover - only runs on regression + raise AssertionError('IMF inference should not use DDIM scheduler step') + + +def _make_images(batch_size, obs_horizon, per_camera_fill): + return { + name: torch.full((batch_size, obs_horizon, 1, 2, 2), fill_value=value, dtype=torch.float32) + for name, value in per_camera_fill.items() + } + + +class IMFVLAAgentTest(unittest.TestCase): + def _make_agent(self, pred_horizon=3, obs_horizon=2, num_action_steps=2): + agent_cls, agent_module = _load_imf_agent_class() + head = _RecordingLinearIMFHead() + agent = agent_cls( + vision_backbone=_StubVisionBackbone(), + state_encoder=nn.Identity(), + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=pred_horizon, + obs_horizon=obs_horizon, + diffusion_steps=10, + inference_steps=1, + num_cams=len(_CAMERA_NAMES), + camera_names=_CAMERA_NAMES, + num_action_steps=num_action_steps, + head_type='transformer', + ) + return agent, head, agent_module + + def test_compute_loss_matches_imf_objective_and_masks_padded_actions(self): + agent, head, agent_module = self._make_agent(pred_horizon=3, obs_horizon=2) + images = _make_images( + batch_size=1, + obs_horizon=2, + per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0}, + ) + qpos = torch.tensor([[[0.25], [0.75]]], dtype=torch.float32) + actions = torch.tensor( + [[[1.0, -1.0], [0.5, 0.25], [-0.5, 1.5]]], + dtype=torch.float32, + ) + action_is_pad = torch.tensor([[False, False, True]]) + noise = torch.tensor( + [[[0.2, -0.4], [0.1, 0.3], [0.5, -0.2]]], + dtype=torch.float32, + ) + t_sample = torch.tensor([0.8], dtype=torch.float32) + r_sample = torch.tensor([0.25], dtype=torch.float32) + + with mock.patch.object(agent_module.torch, 'randn_like', return_value=noise), \ + mock.patch.object(agent_module.torch, 'rand', side_effect=[t_sample, r_sample]): + loss = agent.compute_loss( + { + 'images': images, + 'qpos': qpos, + 'action': actions, + 'action_is_pad': action_is_pad, + } + ) + + cond = torch.tensor([[[1.0, 2.0, 3.0, 0.25], [1.0, 2.0, 3.0, 0.75]]], dtype=torch.float32) + cond_term = cond.mean(dim=(1, 2), keepdim=True) + t = t_sample + r = r_sample + z_t = (1 - t.view(1, 1, 1)) * actions + t.view(1, 1, 1) * noise + scale = head.scale.detach() + u = scale * z_t + r.view(1, 1, 1) + 2.0 * t.view(1, 1, 1) + cond_term + v = scale * z_t + 3.0 * t.view(1, 1, 1) + cond_term + du_dt = scale * v + 2.0 + compound_velocity = u + (t - r).view(1, 1, 1) * du_dt + target = noise - actions + elementwise_loss = (compound_velocity - target) ** 2 + mask = (~action_is_pad).unsqueeze(-1).to(elementwise_loss.dtype) + expected_loss = (elementwise_loss * mask).sum() / (mask.sum() * elementwise_loss.shape[-1]) + + self.assertAlmostEqual(loss.item(), expected_loss.item(), places=6) + self.assertEqual(len(head.calls), 2) + self.assertTrue(torch.allclose(head.calls[0]['r'], t_sample)) + self.assertTrue(torch.allclose(head.calls[0]['t'], t_sample)) + self.assertTrue(torch.allclose(head.calls[0]['cond'], cond)) + + def test_predict_action_uses_one_step_imf_sampling_and_image_conditioning(self): + agent, head, agent_module = self._make_agent(pred_horizon=3, obs_horizon=2) + agent.infer_scheduler = _ForbiddenScheduler() + + images = _make_images( + batch_size=2, + obs_horizon=2, + per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0}, + ) + qpos = torch.tensor( + [ + [[1.0], [2.0]], + [[3.0], [4.0]], + ], + dtype=torch.float32, + ) + initial_noise = torch.tensor( + [ + [[1.0, -1.0], [0.0, 2.0], [3.0, -2.0]], + [[-1.0, 1.0], [2.0, -3.0], [0.5, 0.25]], + ], + dtype=torch.float32, + ) + + with mock.patch.object(agent_module.torch, 'randn', return_value=initial_noise): + predicted_actions = agent.predict_action(images, qpos) + + expected_cond = torch.tensor( + [ + [[10.0, 20.0, 30.0, 1.0], [10.0, 20.0, 30.0, 2.0]], + [[10.0, 20.0, 30.0, 3.0], [10.0, 20.0, 30.0, 4.0]], + ], + dtype=torch.float32, + ) + cond_term = expected_cond.mean(dim=(1, 2), keepdim=True) + expected_actions = 0.5 * initial_noise - 2.0 - cond_term + + self.assertEqual(predicted_actions.shape, (2, 3, 2)) + self.assertTrue(torch.allclose(predicted_actions, expected_actions)) + self.assertEqual(len(head.calls), 1) + self.assertTrue(torch.allclose(head.calls[0]['r'], torch.zeros(2))) + self.assertTrue(torch.allclose(head.calls[0]['t'], torch.ones(2))) + self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond)) + + def test_select_action_only_regenerates_when_action_queue_is_empty(self): + agent, _head, _agent_module = self._make_agent(pred_horizon=4, obs_horizon=2, num_action_steps=2) + observation = { + 'qpos': torch.tensor([0.25], dtype=torch.float32), + 'images': { + 'front': torch.full((1, 2, 2), 3.0, dtype=torch.float32), + 'top': torch.full((1, 2, 2), 2.0, dtype=torch.float32), + 'r_vis': torch.full((1, 2, 2), 1.0, dtype=torch.float32), + }, + } + first_chunk = torch.tensor( + [[[10.0, 11.0], [12.0, 13.0], [14.0, 15.0], [16.0, 17.0]]], + dtype=torch.float32, + ) + second_chunk = torch.tensor( + [[[20.0, 21.0], [22.0, 23.0], [24.0, 25.0], [26.0, 27.0]]], + dtype=torch.float32, + ) + + with mock.patch.object(agent, 'predict_action_chunk', side_effect=[first_chunk, second_chunk]) as mock_predict_chunk: + first_action = agent.select_action(observation) + second_action = agent.select_action(observation) + third_action = agent.select_action(observation) + + self.assertTrue(torch.equal(first_action, first_chunk[0, 1])) + self.assertTrue(torch.equal(second_action, first_chunk[0, 2])) + self.assertTrue(torch.equal(third_action, second_chunk[0, 1])) + self.assertEqual(mock_predict_chunk.call_count, 2) + + def test_joint_visual_backbone_uses_joint_output_dim_for_conditioning(self): + agent_cls, _agent_module = _load_imf_agent_class() + head = _RecordingLinearIMFHead() + vision_backbone = _StubJointVisionBackbone() + agent = agent_cls( + vision_backbone=vision_backbone, + state_encoder=nn.Identity(), + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=3, + obs_horizon=2, + diffusion_steps=10, + inference_steps=1, + num_cams=len(_CAMERA_NAMES), + camera_names=_CAMERA_NAMES, + num_action_steps=2, + head_type='transformer', + ) + + self.assertEqual(agent.per_step_cond_dim, vision_backbone.joint_output_dim + agent.obs_dim) + self.assertEqual( + agent.global_cond_dim, + vision_backbone.joint_output_dim * agent.obs_horizon + agent.obs_dim * agent.obs_horizon, + ) + + images = _make_images( + batch_size=1, + obs_horizon=2, + per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0}, + ) + qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32) + initial_noise = torch.tensor( + [[[1.0, -1.0], [0.0, 2.0], [3.0, -2.0]]], + dtype=torch.float32, + ) + + with mock.patch.object(torch, 'randn', return_value=initial_noise): + predicted_actions = agent.predict_action(images, qpos) + + self.assertEqual(predicted_actions.shape, (1, 3, 2)) + self.assertEqual(len(head.calls), 1) + expected_cond = torch.tensor( + [[[30.0, 20.0, 10.0, 50.0, -20.0, 1.0], [30.0, 20.0, 10.0, 50.0, -19.0, 2.0]]], + dtype=torch.float32, + ) + self.assertEqual(head.calls[0]['cond'].shape[-1], 6) + self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond)) + + def test_multitoken_visual_backbone_flattens_camera_tokens_and_projects_each_with_state(self): + agent_cls, _agent_module = _load_imf_agent_class() + head = _RecordingLinearIMFHead() + projector = nn.Linear(3, 4, bias=False) + with torch.no_grad(): + projector.weight.copy_( + torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 0.0, 1.0], + ], + dtype=torch.float32, + ) + ) + agent = agent_cls( + vision_backbone=_StubMultiTokenVisionBackbone(), + state_encoder=nn.Identity(), + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=3, + obs_horizon=2, + diffusion_steps=10, + inference_steps=1, + num_cams=len(_CAMERA_NAMES), + camera_names=_CAMERA_NAMES, + num_action_steps=2, + head_type='transformer', + cond_projector=projector, + ) + + self.assertEqual(agent.condition_tokens_per_step, 3) + self.assertEqual(agent.condition_sequence_length, 6) + self.assertEqual(agent.per_step_cond_dim, 4) + self.assertEqual(agent.global_cond_dim, 24) + + images = _make_images( + batch_size=1, + obs_horizon=2, + per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0}, + ) + qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32) + cond = agent._build_cond(images, qpos) + + expected = torch.tensor( + [ + [ + [10.0, 10.5, 1.0, 11.0], + [20.0, 20.5, 1.0, 21.0], + [30.0, 30.5, 1.0, 31.0], + [10.0, 10.5, 2.0, 12.0], + [20.0, 20.5, 2.0, 22.0], + [30.0, 30.5, 2.0, 32.0], + ] + ], + dtype=torch.float32, + ) + self.assertEqual(cond.shape, (1, 6, 4)) + self.assertTrue(torch.allclose(cond, expected)) + + def test_multi_token_visual_backbone_pairs_state_per_camera_and_flattens_condition_sequence(self): + agent_cls, agent_module = _load_imf_agent_class() + head = _RecordingLinearIMFHead() + cond_projector = nn.Linear(3, 4, bias=False) + with torch.no_grad(): + cond_projector.weight.copy_(torch.tensor([ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 0.0, 1.0], + ], dtype=torch.float32)) + + agent = agent_cls( + vision_backbone=_StubMultiTokenVisionBackbone(), + state_encoder=nn.Identity(), + action_encoder=nn.Identity(), + head=head, + action_dim=2, + obs_dim=1, + pred_horizon=3, + obs_horizon=2, + diffusion_steps=10, + inference_steps=1, + num_cams=len(_CAMERA_NAMES), + camera_names=_CAMERA_NAMES, + num_action_steps=2, + head_type='transformer', + cond_projector=cond_projector, + ) + agent.infer_scheduler = _ForbiddenScheduler() + + images = _make_images( + batch_size=1, + obs_horizon=2, + per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0}, + ) + qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32) + initial_noise = torch.tensor([[[1.0, -1.0], [0.0, 2.0], [3.0, -2.0]]], dtype=torch.float32) + + with mock.patch.object(agent_module.torch, 'randn', return_value=initial_noise): + predicted_actions = agent.predict_action(images, qpos) + + expected_cond = torch.tensor([[[10.0, 10.5, 1.0, 11.0], + [20.0, 20.5, 1.0, 21.0], + [30.0, 30.5, 1.0, 31.0], + [10.0, 10.5, 2.0, 12.0], + [20.0, 20.5, 2.0, 22.0], + [30.0, 30.5, 2.0, 32.0]]], dtype=torch.float32) + + self.assertEqual(agent.condition_tokens_per_step, 3) + self.assertEqual(agent.condition_sequence_length, 6) + self.assertEqual(agent.raw_per_step_cond_dim, 3) + self.assertEqual(agent.per_step_cond_dim, 4) + self.assertEqual(agent.global_cond_dim, 24) + self.assertEqual(predicted_actions.shape, (1, 3, 2)) + self.assertEqual(len(head.calls), 1) + self.assertEqual(head.calls[0]['cond'].shape, (1, 6, 4)) + self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond)) + + def test_hydra_config_instantiates_resnet_imf_attnres_with_stub_head(self): + cfg = _compose_cfg( + overrides=[ + 'agent=resnet_imf_attnres', + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + 'agent.vision_backbone.freeze_backbone=false', + 'agent.head.n_layer=1', + 'agent.head.n_emb=16', + ] + ) + + self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent') + self.assertEqual(cfg.agent.head._target_, 'roboimi.vla.models.heads.imf_transformer1d.IMFTransformer1D') + self.assertEqual(cfg.agent.head.backbone_type, 'attnres_full') + self.assertEqual(cfg.agent.head.n_head, 1) + self.assertEqual(cfg.agent.head.n_kv_head, 1) + self.assertEqual(cfg.agent.head.n_cond_layers, 0) + self.assertTrue(cfg.agent.head.time_as_cond) + self.assertFalse(cfg.agent.head.causal_attn) + self.assertEqual(cfg.agent.inference_steps, 1) + self.assertEqual(list(cfg.agent.camera_names), list(_CAMERA_NAMES)) + + with _stub_optional_modules(include_imf_head=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.head_type, 'transformer') + self.assertEqual(agent.per_step_cond_dim, agent.vision_encoder.output_dim * agent.num_cams + agent.obs_dim) + self.assertIsInstance(agent.noise_pred_net, _StubIMFHead) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], agent.per_step_cond_dim) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['backbone_type'], 'attnres_full') + + def test_hydra_config_instantiates_resnet_imf_attnres_with_full_attnres_vision_backbone(self): + cfg = _compose_cfg( + overrides=[ + 'agent=resnet_imf_attnres', + 'agent.vision_backbone.vision_backbone_mode=attnres_resnet', + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,56,56]', + 'agent.vision_backbone.freeze_backbone=false', + 'agent.vision_backbone.attnres_stem_dim=16', + 'agent.vision_backbone.attnres_stage_dims=[16,32,64,128]', + 'agent.vision_backbone.attnres_stage_depths=[1,1,1,1]', + 'agent.vision_backbone.attnres_stage_heads=[2,4,4,8]', + 'agent.vision_backbone.attnres_stage_kv_heads=[1,1,1,1]', + 'agent.vision_backbone.attnres_stage_window_sizes=[7,7,7,7]', + 'agent.head.n_layer=1', + 'agent.head.n_emb=16', + ] + ) + + with _stub_optional_modules(include_imf_head=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.vision_encoder.output_dim, 64) + self.assertEqual(agent.per_step_cond_dim, 64 * agent.num_cams + agent.obs_dim) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], agent.per_step_cond_dim) + + def test_hydra_config_instantiates_lewm_imf_attnres_with_joint_visual_condition_dim(self): + cfg = _compose_cfg( + overrides=[ + 'agent=lewm_imf_attnres', + 'agent.vision_backbone.checkpoint_path=null', + 'agent.head.n_layer=1', + 'agent.head.n_emb=16', + ] + ) + + self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent') + self.assertEqual(cfg.agent.vision_backbone._target_, 'roboimi.vla.models.backbones.lewm_vit_backbone.LEWMViTBackbone') + self.assertEqual(list(cfg.agent.camera_names), list(_CAMERA_NAMES)) + self.assertEqual(list(cfg.agent.vision_backbone.camera_names), list(_CAMERA_NAMES)) + self.assertEqual(list(cfg.agent.vision_backbone.fused_camera_names), ['front', 'top', 'r_vis']) + self.assertIsNone(cfg.agent.vision_backbone.dataset_image_resize_shape) + self.assertEqual(list(cfg.agent.vision_backbone.eval_image_resize_shape), [256, 256]) + self.assertEqual(cfg.agent.head.cond_dim, 208) + + with _stub_optional_modules(include_imf_head=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.per_step_cond_dim, agent.vision_encoder.joint_output_dim + agent.obs_dim) + self.assertEqual(agent.per_step_cond_dim, 208) + self.assertEqual(agent.global_cond_dim, agent.obs_horizon * 208) + self.assertIsNone(agent.vision_encoder.dataset_image_resize_shape) + self.assertEqual(agent.vision_encoder.eval_image_resize_shape, (256, 256)) + self.assertIsInstance(agent.noise_pred_net, _StubIMFHead) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 208) + + def test_hydra_config_instantiates_resnet_imf_attnres_multitoken_with_projected_camera_tokens(self): + cfg = _compose_cfg( + overrides=[ + 'agent=resnet_imf_attnres_multitoken', + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + 'agent.head.n_layer=1', + 'agent.head.n_emb=32', + ] + ) + + self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent') + self.assertEqual(cfg.agent.vision_backbone.vision_backbone_mode, 'resnet') + self.assertTrue(cfg.agent.vision_backbone.use_separate_rgb_encoder_per_camera) + self.assertTrue(cfg.agent.vision_backbone.output_tokens_per_camera) + self.assertEqual(cfg.agent.cond_projector.output_dim, 32) + self.assertEqual(cfg.agent.head.cond_dim, 32) + + with _stub_optional_modules(include_imf_head=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.condition_tokens_per_step, 3) + self.assertEqual(agent.condition_sequence_length, agent.obs_horizon * 3) + self.assertEqual(agent.per_step_cond_dim, 32) + self.assertEqual(agent.global_cond_dim, agent.condition_sequence_length * 32) + self.assertIsInstance(agent.noise_pred_net, _StubIMFHead) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 32) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], 6) + + + def test_hydra_config_instantiates_siglip2_imf_attnres_with_condition_projection(self): + cfg = _compose_cfg( + overrides=[ + 'agent=siglip2_imf_attnres', + 'agent.vision_backbone.per_view_output_dim=96', + 'agent.head.n_layer=1', + 'agent.head.n_emb=16', + 'agent.cond_projector.output_dim=384', + ] + ) + + self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent') + self.assertEqual( + cfg.agent.vision_backbone._target_, + 'roboimi.vla.models.backbones.siglip2_diffusion_backbone.SigLIP2DiffusionBackbone', + ) + self.assertEqual(list(cfg.agent.camera_names), list(_CAMERA_NAMES)) + self.assertIsNone(cfg.agent.vision_backbone.dataset_image_resize_shape) + self.assertEqual(list(cfg.agent.vision_backbone.eval_image_resize_shape), [256, 256]) + self.assertEqual(cfg.agent.head.cond_dim, 384) + + with _stub_optional_modules(include_imf_head=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.raw_per_step_cond_dim, 3 * 96 + agent.obs_dim) + self.assertEqual(agent.per_step_cond_dim, 384) + self.assertEqual(agent.global_cond_dim, agent.obs_horizon * 384) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 384) + self.assertEqual(agent.vision_encoder.output_dim, 96) + self.assertEqual(agent.vision_encoder.eval_image_resize_shape, (256, 256)) + + + def test_hydra_config_instantiates_resnet_imf_attnres_multitoken_with_sequence_length_three_times_obs_horizon(self): + cfg = _compose_cfg( + overrides=[ + 'agent=resnet_imf_attnres_multitoken', + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + 'agent.vision_backbone.freeze_backbone=false', + 'agent.head.n_layer=1', + 'agent.head.n_emb=16', + ] + ) + + self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent') + self.assertEqual(list(cfg.agent.camera_names), list(_CAMERA_NAMES)) + self.assertTrue(cfg.agent.vision_backbone.use_separate_rgb_encoder_per_camera) + self.assertTrue(cfg.agent.vision_backbone.output_tokens_per_camera) + self.assertEqual(cfg.agent.vision_backbone.vision_backbone_mode, 'resnet') + self.assertEqual(cfg.agent.cond_projector.output_dim, 16) + self.assertEqual(cfg.agent.head.cond_dim, 16) + + with _stub_optional_modules(include_imf_head=True): + agent = instantiate(cfg.agent) + + self.assertEqual(agent.condition_tokens_per_step, 3) + self.assertEqual(agent.condition_sequence_length, agent.obs_horizon * 3) + self.assertEqual(agent.per_step_cond_dim, 16) + self.assertEqual(agent.global_cond_dim, agent.condition_sequence_length * 16) + self.assertEqual(agent.vision_encoder.tokens_per_step, 3) + self.assertIsInstance(agent.noise_pred_net, _StubIMFHead) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 16) + self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], agent.condition_sequence_length) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_lewm_vit_backbone.py b/tests/test_lewm_vit_backbone.py new file mode 100644 index 0000000..85f06d9 --- /dev/null +++ b/tests/test_lewm_vit_backbone.py @@ -0,0 +1,220 @@ +import tempfile +import types +import unittest +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import ViTConfig, ViTModel + + +_INPUT_CAMERA_NAMES = ("r_vis", "top", "front") +_FUSED_CAMERA_NAMES = ("front", "top", "r_vis") + + +class _ReferenceProjector(nn.Module): + def __init__(self): + super().__init__() + self.net = nn.Sequential( + nn.Linear(192, 2048), + nn.BatchNorm1d(2048), + nn.GELU(), + nn.Linear(2048, 192), + ) + + def forward(self, x): + return self.net(x) + + +def _build_reference_encoder() -> ViTModel: + return ViTModel( + ViTConfig( + image_size=224, + patch_size=14, + num_channels=3, + hidden_size=192, + intermediate_size=768, + num_hidden_layers=12, + num_attention_heads=3, + qkv_bias=True, + ), + add_pooling_layer=False, + ) + + +def _write_synthetic_lightning_ckpt(path: Path): + torch.manual_seed(7) + encoder = _build_reference_encoder() + projector = _ReferenceProjector() + lightning_state_dict = {} + for key, value in encoder.state_dict().items(): + lightning_state_dict[f"model.encoder.{key}"] = value.detach().clone() + for key, value in projector.state_dict().items(): + lightning_state_dict[f"model.projector.{key}"] = value.detach().clone() + torch.save({"state_dict": lightning_state_dict}, path) + return encoder.state_dict(), projector.state_dict() + + +class LEWMViTBackboneTest(unittest.TestCase): + def test_loads_lightning_encoder_and_projector_checkpoint_and_emits_joint_embedding(self): + from roboimi.vla.models.backbones.lewm_vit_backbone import LEWMViTBackbone + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt_path = Path(tmpdir) / "synthetic-lewm.ckpt" + reference_encoder_state, reference_projector_state = _write_synthetic_lightning_ckpt( + ckpt_path + ) + + backbone = LEWMViTBackbone( + checkpoint_path=ckpt_path, + camera_names=_INPUT_CAMERA_NAMES, + fused_camera_names=_FUSED_CAMERA_NAMES, + freeze_backbone=True, + ) + + self.assertEqual(backbone.camera_names, _INPUT_CAMERA_NAMES) + self.assertEqual(backbone.fused_camera_names, _FUSED_CAMERA_NAMES) + self.assertEqual(backbone.num_cameras, 3) + self.assertEqual(backbone.joint_output_dim, 192) + self.assertEqual(backbone.output_dim, 192) + self.assertEqual(backbone.encoder.config.hidden_size, 192) + self.assertEqual(backbone.encoder.config.patch_size, 14) + self.assertEqual(backbone.encoder.config.num_hidden_layers, 12) + self.assertEqual(backbone.encoder.config.num_attention_heads, 3) + + for key, value in reference_encoder_state.items(): + self.assertTrue(torch.equal(backbone.encoder.state_dict()[key], value), key) + for key, value in reference_projector_state.items(): + self.assertTrue(torch.equal(backbone.projector.state_dict()[key], value), key) + + images = { + cam_name: torch.rand(1, 1, 3, 224, 224) + for cam_name in _INPUT_CAMERA_NAMES + } + output = backbone(images) + + self.assertEqual(output.shape, (1, 1, 192)) + self.assertFalse(output.requires_grad) + + def test_forward_uses_front_top_rvis_fusion_order_and_exact_lewm_cwh_resize_path(self): + from roboimi.vla.models.backbones.lewm_vit_backbone import LEWMViTBackbone + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt_path = Path(tmpdir) / "synthetic-lewm.ckpt" + _write_synthetic_lightning_ckpt(ckpt_path) + + backbone = LEWMViTBackbone( + checkpoint_path=ckpt_path, + camera_names=_INPUT_CAMERA_NAMES, + fused_camera_names=_FUSED_CAMERA_NAMES, + freeze_backbone=True, + ) + captured = {} + + def fake_encoder_forward(module, pixel_values, interpolate_pos_encoding=False, **kwargs): + del module, kwargs + captured["pixel_values"] = pixel_values.detach().clone() + captured["interpolate_pos_encoding"] = interpolate_pos_encoding + batch = pixel_values.shape[0] + patch_tokens = (pixel_values.shape[-2] // 14) * (pixel_values.shape[-1] // 14) + cls = ( + torch.arange(192, dtype=pixel_values.dtype, device=pixel_values.device) + .unsqueeze(0) + .expand(batch, -1) + ) + last_hidden_state = torch.zeros( + batch, + patch_tokens + 1, + 192, + dtype=pixel_values.dtype, + device=pixel_values.device, + ) + last_hidden_state[:, 0] = cls + return types.SimpleNamespace(last_hidden_state=last_hidden_state) + + backbone.encoder.forward = types.MethodType(fake_encoder_forward, backbone.encoder) + + r_vis = torch.full((1, 1, 3, 256, 256), 0.30) + top = torch.full((1, 1, 3, 256, 256), 0.20) + front = torch.full((1, 1, 3, 256, 256), 0.10) + bn = backbone.projector.net[1] + running_mean_before = bn.running_mean.detach().clone() + running_var_before = bn.running_var.detach().clone() + + backbone.train() + self.assertFalse(backbone.encoder.training) + self.assertFalse(backbone.projector.training) + + output = backbone({"r_vis": r_vis, "top": top, "front": front}) + + self.assertEqual(output.shape, (1, 1, 192)) + self.assertEqual(captured["pixel_values"].shape, (1, 3, 672, 224)) + self.assertTrue(captured["interpolate_pos_encoding"]) + + normalized_views = [ + ((view.reshape(-1, *view.shape[2:]).float()).clamp(0.0, 1.0) - backbone.mean) / backbone.std + for view in (front, top, r_vis) + ] + expected_fuse_then_resize = F.interpolate( + torch.cat(normalized_views, dim=-2), + size=(672, 224), + mode="bilinear", + align_corners=False, + antialias=True, + ) + expected_pre_resize_then_fuse = torch.cat( + [ + F.interpolate( + view, + size=(224, 224), + mode="bilinear", + align_corners=False, + antialias=True, + ) + for view in normalized_views + ], + dim=-2, + ) + + self.assertTrue( + torch.allclose(captured["pixel_values"], expected_fuse_then_resize, atol=1e-6, rtol=1e-6) + ) + self.assertFalse( + torch.allclose( + expected_fuse_then_resize, + expected_pre_resize_then_fuse, + atol=1e-6, + rtol=1e-6, + ) + ) + self.assertFalse( + torch.allclose( + captured["pixel_values"], + expected_pre_resize_then_fuse, + atol=1e-6, + rtol=1e-6, + ) + ) + self.assertTrue( + torch.allclose( + captured["pixel_values"][0, :, 223, :], + expected_fuse_then_resize[0, :, 223, :], + atol=1e-6, + rtol=1e-6, + ) + ) + self.assertTrue( + torch.allclose( + captured["pixel_values"][0, :, 447, :], + expected_fuse_then_resize[0, :, 447, :], + atol=1e-6, + rtol=1e-6, + ) + ) + self.assertTrue(torch.equal(bn.running_mean, running_mean_before)) + self.assertTrue(torch.equal(bn.running_var, running_var_before)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_raw_action_trajectory_viewer.py b/tests/test_raw_action_trajectory_viewer.py new file mode 100644 index 0000000..8a15524 --- /dev/null +++ b/tests/test_raw_action_trajectory_viewer.py @@ -0,0 +1,119 @@ +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import numpy as np + +from roboimi.utils import raw_action_trajectory_viewer as traj_view + + +class RawActionTrajectoryViewerTest(unittest.TestCase): + def test_set_transfer_box_pose_writes_joint_qpos(self): + joint_qpos = np.zeros(7, dtype=np.float64) + + class _FakeJoint: + def __init__(self, qpos): + self.qpos = qpos + + class _FakeData: + def joint(self, name): + assert name == "red_box_joint" + return _FakeJoint(joint_qpos) + + traj_view.set_transfer_box_pose(_FakeData(), np.array([0.2, -0.1, 1.05], dtype=np.float64)) + + np.testing.assert_array_equal( + joint_qpos, + np.array([0.2, -0.1, 1.05, 1.0, 0.0, 0.0, 0.0], dtype=np.float64), + ) + + def test_disable_cv2_highgui_temporarily_replaces_gui_calls(self): + fake_cv2 = SimpleNamespace( + namedWindow=lambda *args, **kwargs: "named", + imshow=lambda *args, **kwargs: "imshow", + waitKey=lambda *args, **kwargs: "wait", + ) + + restore = traj_view.disable_cv2_highgui(fake_cv2) + self.assertIsNone(fake_cv2.namedWindow("x")) + self.assertIsNone(fake_cv2.imshow("x", None)) + self.assertEqual(fake_cv2.waitKey(1), 1) + + restore() + self.assertEqual(fake_cv2.namedWindow("x"), "named") + self.assertEqual(fake_cv2.imshow("x", None), "imshow") + self.assertEqual(fake_cv2.waitKey(1), "wait") + + def test_load_raw_action_positions_from_npz(self): + raw_action = np.array( + [ + [1.0, 2.0, 3.0, 0, 0, 0, 1, 11.0, 12.0, 13.0, 0, 0, 0, 1, -1, -1], + [4.0, 5.0, 6.0, 0, 0, 0, 1, 14.0, 15.0, 16.0, 0, 0, 0, 1, -1, -1], + ], + dtype=np.float32, + ) + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "trajectory.npz" + np.savez(path, raw_action=raw_action) + + positions = traj_view.load_raw_action_positions(path) + + np.testing.assert_array_equal( + positions["left"], + np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32), + ) + np.testing.assert_array_equal( + positions["right"], + np.array([[11.0, 12.0, 13.0], [14.0, 15.0, 16.0]], dtype=np.float32), + ) + + def test_build_red_capsule_segments_downsamples_to_fit_scene_limit(self): + left = np.stack([np.array([float(i), 0.0, 0.0], dtype=np.float32) for i in range(6)]) + right = np.stack([np.array([float(i), 1.0, 0.0], dtype=np.float32) for i in range(6)]) + + markers = traj_view.build_trajectory_capsule_markers( + {"left": left, "right": right}, + max_markers=4, + radius=0.01, + ) + + self.assertLessEqual(len(markers), 4) + self.assertTrue(all(marker["rgba"] == (1.0, 0.0, 0.0, 1.0) for marker in markers)) + self.assertTrue(all(marker["radius"] == 0.01 for marker in markers)) + + def test_apply_capsule_markers_populates_user_scene(self): + fake_scene = SimpleNamespace( + maxgeom=3, + ngeom=99, + geoms=[object(), object(), object()], + ) + markers = [ + { + "from": np.array([0.0, 0.0, 0.0], dtype=np.float64), + "to": np.array([1.0, 0.0, 0.0], dtype=np.float64), + "rgba": (1.0, 0.0, 0.0, 1.0), + "radius": 0.01, + }, + { + "from": np.array([0.0, 1.0, 0.0], dtype=np.float64), + "to": np.array([1.0, 1.0, 0.0], dtype=np.float64), + "rgba": (1.0, 0.0, 0.0, 1.0), + "radius": 0.01, + }, + ] + + with mock.patch.object(traj_view.mujoco, "mjv_initGeom") as init_geom, mock.patch.object( + traj_view.mujoco, + "mjv_connector", + ) as connector: + traj_view.apply_capsule_markers_to_scene(fake_scene, markers) + + self.assertEqual(fake_scene.ngeom, 2) + self.assertEqual(init_geom.call_count, 2) + self.assertEqual(connector.call_count, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_resnet_transformer_agent_wiring.py b/tests/test_resnet_transformer_agent_wiring.py new file mode 100644 index 0000000..c9a0ef2 --- /dev/null +++ b/tests/test_resnet_transformer_agent_wiring.py @@ -0,0 +1,455 @@ +import contextlib +import sys +import types +import unittest +from pathlib import Path + +import torch +from hydra import compose, initialize_config_dir +from hydra.errors import InstantiationException +from hydra.core.global_hydra import GlobalHydra +from hydra.utils import instantiate +from omegaconf import OmegaConf + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve()) +_EXPECTED_CAMERA_NAMES = ['r_vis', 'top', 'front'] +_MISSING = object() + + +class _FakeScheduler: + def __init__(self, num_train_timesteps=100, **kwargs): + self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps) + self.timesteps = [] + + def add_noise(self, sample, noise, timestep): + return sample + noise + + def set_timesteps(self, num_inference_steps): + self.timesteps = list(range(num_inference_steps - 1, -1, -1)) + + def step(self, noise_pred, timestep, sample): + return types.SimpleNamespace(prev_sample=sample) + + +class _IdentityCrop: + def __init__(self, size): + self.size = size + + def __call__(self, x): + return x + + +class _FakeResNet(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu1 = torch.nn.ReLU() + self.conv2 = torch.nn.Conv2d(8, 16, kernel_size=3, padding=1, stride=2) + self.relu2 = torch.nn.ReLU() + self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1)) + self.fc = torch.nn.Linear(16, 16) + + def forward(self, x): + x = self.relu1(self.conv1(x)) + x = self.relu2(self.conv2(x)) + x = self.avgpool(x) + x = torch.flatten(x, start_dim=1) + return self.fc(x) + + +class _FakeRearrange(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, x): + return x + + +class _CondCapturingHead(torch.nn.Module): + def __init__(self): + super().__init__() + self.last_cond = None + + def forward(self, sample, timestep, cond): + self.last_cond = cond.detach().clone() + return torch.zeros_like(sample) + + +@contextlib.contextmanager +def _stub_optional_modules(): + previous_modules = {} + + def inject(name, module): + if name not in previous_modules: + previous_modules[name] = sys.modules.get(name, _MISSING) + sys.modules[name] = module + + diffusers_module = types.ModuleType('diffusers') + schedulers_module = types.ModuleType('diffusers.schedulers') + ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm') + ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim') + ddpm_module.DDPMScheduler = _FakeScheduler + ddim_module.DDIMScheduler = _FakeScheduler + diffusers_module.DDPMScheduler = _FakeScheduler + diffusers_module.DDIMScheduler = _FakeScheduler + diffusers_module.schedulers = schedulers_module + schedulers_module.scheduling_ddpm = ddpm_module + schedulers_module.scheduling_ddim = ddim_module + + torchvision_module = types.ModuleType('torchvision') + models_module = types.ModuleType('torchvision.models') + transforms_module = types.ModuleType('torchvision.transforms') + models_module.resnet18 = lambda weights=None: _FakeResNet() + transforms_module.CenterCrop = _IdentityCrop + transforms_module.RandomCrop = _IdentityCrop + torchvision_module.models = models_module + torchvision_module.transforms = transforms_module + + einops_module = types.ModuleType('einops') + einops_module.rearrange = lambda x, *args, **kwargs: x + einops_layers_module = types.ModuleType('einops.layers') + einops_layers_torch_module = types.ModuleType('einops.layers.torch') + einops_layers_torch_module.Rearrange = _FakeRearrange + einops_module.layers = einops_layers_module + einops_layers_module.torch = einops_layers_torch_module + + try: + inject('diffusers', diffusers_module) + inject('diffusers.schedulers', schedulers_module) + inject('diffusers.schedulers.scheduling_ddpm', ddpm_module) + inject('diffusers.schedulers.scheduling_ddim', ddim_module) + inject('torchvision', torchvision_module) + inject('torchvision.models', models_module) + inject('torchvision.transforms', transforms_module) + inject('einops', einops_module) + inject('einops.layers', einops_layers_module) + inject('einops.layers.torch', einops_layers_torch_module) + yield + finally: + for name, previous in reversed(list(previous_modules.items())): + if previous is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + +def _compose_cfg(overrides=None): + if not OmegaConf.has_resolver('len'): + OmegaConf.register_new_resolver('len', lambda x: len(x)) + + GlobalHydra.instance().clear() + with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR): + return compose(config_name='config', overrides=list(overrides or [])) + + +def _make_images(batch_size, obs_horizon, image_shape, per_camera_fill=None): + channels, height, width = image_shape + per_camera_fill = per_camera_fill or { + 'front': 30.0, + 'top': 20.0, + 'r_vis': 10.0, + } + return { + name: torch.full( + (batch_size, obs_horizon, channels, height, width), + fill_value=fill_value, + dtype=torch.float32, + ) + for name, fill_value in per_camera_fill.items() + } + + +def _patch_backbone_for_order_tracking(backbone): + feature_dim = backbone.output_dim + + def encode_mean(image_batch): + mean_feature = image_batch.mean(dim=(1, 2, 3)).unsqueeze(-1) + return mean_feature.repeat(1, feature_dim) + + if backbone.use_separate_rgb_encoder_per_camera: + for encoder in backbone.rgb_encoder: + encoder.forward_single_image = encode_mean + else: + backbone.rgb_encoder.forward_single_image = encode_mean + + +def _extract_camera_markers(cond, feature_dim, num_cams): + camera_block = cond[0, 0, : feature_dim * num_cams].view(num_cams, feature_dim) + return camera_block[:, 0] + + +def _extract_token_camera_markers(tokens): + return tokens[0, 0, :, 0] + + +def _extract_token_markers(token_sequence): + return token_sequence[0, 0, :, 0] + + +class ResNetTransformerAgentWiringTest(unittest.TestCase): + def test_hydra_wiring_uses_required_three_camera_transformer_conditioning_in_agent_order_and_ignores_extra_keys(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + 'agent.inference_steps=1', + 'agent.head.n_layer=1', + 'agent.head.n_cond_layers=0', + 'agent.head.n_emb=32', + 'agent.head.n_head=4', + ] + ) + + self.assertEqual(list(cfg.data.camera_names), _EXPECTED_CAMERA_NAMES) + self.assertEqual(list(cfg.eval.camera_names), _EXPECTED_CAMERA_NAMES) + self.assertEqual(list(cfg.agent.camera_names), _EXPECTED_CAMERA_NAMES) + self.assertEqual(list(cfg.agent.vision_backbone.camera_names), _EXPECTED_CAMERA_NAMES) + self.assertEqual(cfg.agent.head_type, 'transformer') + self.assertEqual(cfg.agent.num_cams, 3) + self.assertTrue(cfg.agent.head.obs_as_cond) + self.assertFalse(cfg.agent.head.causal_attn) + + with _stub_optional_modules(): + agent = instantiate(cfg.agent) + expected_cond_dim = agent.vision_encoder.output_dim * agent.num_cams + agent.obs_dim + self.assertEqual(cfg.agent.head.cond_dim, expected_cond_dim) + self.assertEqual(agent.per_step_cond_dim, expected_cond_dim) + self.assertEqual(agent.noise_pred_net.cond_obs_emb.in_features, expected_cond_dim) + + batch_size = 2 + image_shape = tuple(cfg.agent.vision_backbone.input_shape) + images = _make_images( + batch_size, + cfg.agent.obs_horizon, + image_shape, + per_camera_fill={ + 'front': 30.0, + 'top': 20.0, + 'r_vis': 10.0, + 'left_wrist': 99.0, + }, + ) + proprioception = torch.randn(batch_size, cfg.agent.obs_horizon, cfg.agent.obs_dim) + _patch_backbone_for_order_tracking(agent.vision_encoder) + capturing_head = _CondCapturingHead() + agent.noise_pred_net = capturing_head + predicted_actions = agent.predict_action(images, proprioception) + self.assertEqual( + predicted_actions.shape, + (batch_size, cfg.agent.pred_horizon, cfg.agent.action_dim), + ) + self.assertIsNotNone(capturing_head.last_cond) + self.assertEqual(capturing_head.last_cond.shape[-1], expected_cond_dim) + camera_markers = _extract_camera_markers( + capturing_head.last_cond, + agent.vision_encoder.output_dim, + agent.num_cams, + ) + self.assertTrue(torch.allclose(camera_markers, torch.tensor([10.0, 20.0, 30.0]))) + + missing_images = dict(images) + missing_images.pop('top') + with self.assertRaisesRegex(ValueError, 'missing=.*top'): + agent.predict_action(missing_images, proprioception) + + def test_multitoken_resnet_backbone_emits_one_token_per_camera_in_agent_order(self): + cfg = _compose_cfg( + overrides=[ + 'agent=resnet_imf_attnres_multitoken', + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + ] + ) + + with _stub_optional_modules(): + backbone = instantiate(cfg.agent.vision_backbone) + _patch_backbone_for_order_tracking(backbone) + images = _make_images( + batch_size=1, + obs_horizon=cfg.agent.obs_horizon, + image_shape=tuple(cfg.agent.vision_backbone.input_shape), + per_camera_fill={ + 'front': 30.0, + 'top': 20.0, + 'r_vis': 10.0, + 'left_wrist': 99.0, + }, + ) + tokens = backbone(images) + + self.assertEqual(tokens.shape, (1, cfg.agent.obs_horizon, 3, backbone.output_dim)) + self.assertEqual(backbone.tokens_per_step, 3) + camera_markers = _extract_token_camera_markers(tokens) + self.assertTrue(torch.allclose(camera_markers, torch.tensor([10.0, 20.0, 30.0]))) + + def test_agent_rejects_conflicting_explicit_backbone_camera_names(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + ] + ) + cfg.agent.vision_backbone.camera_names = ['front', 'top', 'r_vis'] + + with _stub_optional_modules(): + with self.assertRaisesRegex(InstantiationException, 'camera_names'): + instantiate(cfg.agent) + + def test_backbone_uses_sorted_fallback_order_when_camera_names_unset(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + ] + ) + cfg.agent.vision_backbone.camera_names = None + + with _stub_optional_modules(): + backbone = instantiate(cfg.agent.vision_backbone) + _patch_backbone_for_order_tracking(backbone) + images = _make_images( + batch_size=1, + obs_horizon=cfg.agent.obs_horizon, + image_shape=tuple(cfg.agent.vision_backbone.input_shape), + per_camera_fill={ + 'top': 20.0, + 'front': 30.0, + 'r_vis': 10.0, + }, + ) + ordered_features = backbone(images) + camera_markers = _extract_camera_markers( + ordered_features, + backbone.output_dim, + len(images), + ) + self.assertTrue(torch.allclose(camera_markers, torch.tensor([30.0, 10.0, 20.0]))) + + def test_agent_queue_fallback_order_is_deterministic_when_camera_names_unset(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + ] + ) + cfg.agent.camera_names = None + cfg.agent.vision_backbone.camera_names = None + + with _stub_optional_modules(): + agent = instantiate(cfg.agent) + observation = { + 'qpos': torch.randn(cfg.agent.obs_dim), + 'images': { + 'top': torch.full(tuple(cfg.agent.vision_backbone.input_shape), 20.0), + 'front': torch.full(tuple(cfg.agent.vision_backbone.input_shape), 30.0), + 'r_vis': torch.full(tuple(cfg.agent.vision_backbone.input_shape), 10.0), + }, + } + agent._populate_queues(observation) + batch = agent._prepare_observation_batch() + self.assertEqual(list(batch['images'].keys()), ['front', 'r_vis', 'top']) + + def test_backbone_rejects_camera_count_mismatch_when_camera_names_unset(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + ] + ) + cfg.agent.vision_backbone.camera_names = None + + with _stub_optional_modules(): + backbone = instantiate(cfg.agent.vision_backbone) + images = _make_images( + batch_size=1, + obs_horizon=cfg.agent.obs_horizon, + image_shape=tuple(cfg.agent.vision_backbone.input_shape), + per_camera_fill={ + 'front': 30.0, + 'r_vis': 10.0, + }, + ) + with self.assertRaisesRegex(ValueError, 'num_cameras'): + backbone(images) + + def test_agent_rejects_camera_count_mismatch_when_camera_names_unset(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + 'agent.inference_steps=1', + 'agent.head.n_layer=1', + 'agent.head.n_cond_layers=0', + 'agent.head.n_emb=32', + 'agent.head.n_head=4', + ] + ) + cfg.agent.camera_names = None + cfg.agent.vision_backbone.camera_names = None + + with _stub_optional_modules(): + agent = instantiate(cfg.agent) + images = _make_images( + batch_size=1, + obs_horizon=cfg.agent.obs_horizon, + image_shape=tuple(cfg.agent.vision_backbone.input_shape), + per_camera_fill={ + 'front': 30.0, + 'r_vis': 10.0, + }, + ) + proprioception = torch.randn(1, cfg.agent.obs_horizon, cfg.agent.obs_dim) + with self.assertRaisesRegex(ValueError, 'num_cams'): + agent.predict_action(images, proprioception) + + def test_agent_rejects_num_cams_mismatch_with_backbone_when_camera_names_unset(self): + cfg = _compose_cfg( + overrides=[ + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + ] + ) + cfg.agent.camera_names = None + cfg.agent.vision_backbone.camera_names = None + cfg.agent.num_cams = 2 + cfg.agent.vision_backbone.num_cameras = 3 + + with _stub_optional_modules(): + with self.assertRaisesRegex(InstantiationException, 'num_cams'): + instantiate(cfg.agent) + + def test_multitoken_resnet_backbone_emits_one_token_per_camera_in_agent_order(self): + cfg = _compose_cfg( + overrides=[ + 'agent=resnet_imf_attnres_multitoken', + 'agent.vision_backbone.pretrained_backbone_weights=null', + 'agent.vision_backbone.input_shape=[3,16,16]', + 'agent.head.n_layer=1', + 'agent.head.n_emb=32', + ] + ) + + with _stub_optional_modules(): + backbone = instantiate(cfg.agent.vision_backbone) + _patch_backbone_for_order_tracking(backbone) + images = _make_images( + batch_size=1, + obs_horizon=cfg.agent.obs_horizon, + image_shape=tuple(cfg.agent.vision_backbone.input_shape), + per_camera_fill={ + 'front': 30.0, + 'top': 20.0, + 'r_vis': 10.0, + }, + ) + output = backbone(images) + + self.assertEqual(output.shape, (1, cfg.agent.obs_horizon, 3, backbone.output_dim)) + token_markers = _extract_token_markers(output) + self.assertTrue(torch.allclose(token_markers, torch.tensor([10.0, 20.0, 30.0]))) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_robot_asset_paths.py b/tests/test_robot_asset_paths.py new file mode 100644 index 0000000..5c2fd08 --- /dev/null +++ b/tests/test_robot_asset_paths.py @@ -0,0 +1,105 @@ +import os +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from roboimi.assets.robots import diana_med + + +class _FakeKDL: + init_calls = [] + reset_calls = [] + + def __init__(self, urdf_path): + self.__class__.init_calls.append(urdf_path) + + def resetChain(self, base, end): + self.__class__.reset_calls.append((base, end)) + + +class RobotAssetPathResolutionTest(unittest.TestCase): + def setUp(self): + _FakeKDL.init_calls = [] + _FakeKDL.reset_calls = [] + + def test_bidianamed_resolves_robot_asset_paths_independent_of_cwd(self): + BiDianaMed = diana_med.BiDianaMed + repo_root = Path(__file__).resolve().parents[1] + expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml' + expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf' + xml_calls = [] + + def fake_from_xml_path(*, filename, assets=None): + xml_calls.append((filename, assets)) + return object() + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch( + 'roboimi.assets.robots.arm_base.mujoco.MjModel.from_xml_path', + side_effect=fake_from_xml_path, + ), mock.patch( + 'roboimi.assets.robots.arm_base.mujoco.MjData', + return_value=object(), + ), mock.patch( + 'roboimi.assets.robots.arm_base.KDL_utils', + _FakeKDL, + ): + BiDianaMed() + finally: + os.chdir(previous_cwd) + + self.assertEqual(len(xml_calls), 1) + self.assertEqual(Path(xml_calls[0][0]), expected_xml) + self.assertTrue(Path(xml_calls[0][0]).is_absolute()) + self.assertGreaterEqual(len(_FakeKDL.init_calls), 2) + self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf}) + self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls)) + + def test_bidianamed_socket_peg_resolves_robot_asset_paths_independent_of_cwd(self): + BiDianaMedSocketPeg = getattr(diana_med, 'BiDianaMedSocketPeg', None) + self.assertIsNotNone( + BiDianaMedSocketPeg, + 'Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg', + ) + + repo_root = Path(__file__).resolve().parents[1] + expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml' + expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf' + xml_calls = [] + + def fake_from_xml_path(*, filename, assets=None): + xml_calls.append((filename, assets)) + return object() + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch( + 'roboimi.assets.robots.arm_base.mujoco.MjModel.from_xml_path', + side_effect=fake_from_xml_path, + ), mock.patch( + 'roboimi.assets.robots.arm_base.mujoco.MjData', + return_value=object(), + ), mock.patch( + 'roboimi.assets.robots.arm_base.KDL_utils', + _FakeKDL, + ): + BiDianaMedSocketPeg() + finally: + os.chdir(previous_cwd) + + self.assertEqual(len(xml_calls), 1) + self.assertEqual(Path(xml_calls[0][0]), expected_xml) + self.assertTrue(Path(xml_calls[0][0]).is_absolute()) + self.assertGreaterEqual(len(_FakeKDL.init_calls), 2) + self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf}) + self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_siglip2_diffusion_backbone.py b/tests/test_siglip2_diffusion_backbone.py new file mode 100644 index 0000000..a39dc08 --- /dev/null +++ b/tests/test_siglip2_diffusion_backbone.py @@ -0,0 +1,121 @@ +import types +import unittest +from unittest import mock + +import torch +from torch import nn + + +_CAMERA_NAMES = ("r_vis", "top", "front") + + +class _FakeSiglipVisionOutput: + def __init__(self, pooler_output): + self.pooler_output = pooler_output + + +class _FakeSiglipVisionConfig: + def __init__(self, hidden_size=768, image_size=256): + self.hidden_size = hidden_size + self.image_size = image_size + + +class _FakeSiglipVisionModel(nn.Module): + def __init__(self, hidden_size=768): + super().__init__() + self.config = _FakeSiglipVisionConfig(hidden_size=hidden_size) + self.forward_calls = [] + + @classmethod + def from_pretrained(cls, *args, **kwargs): + del args, kwargs + return cls() + + def forward(self, pixel_values=None, **kwargs): + self.forward_calls.append({ + "pixel_values": pixel_values.detach().clone(), + "kwargs": dict(kwargs), + }) + pooled = pixel_values.mean(dim=(2, 3), keepdim=False) + return _FakeSiglipVisionOutput(pooler_output=pooled) + + +class SigLIP2DiffusionBackboneTest(unittest.TestCase): + def test_forward_encodes_each_view_independently_and_concatenates_projected_features(self): + from roboimi.vla.models.backbones.siglip2_diffusion_backbone import SigLIP2DiffusionBackbone + + fake_model = _FakeSiglipVisionModel(hidden_size=3) + with mock.patch( + "roboimi.vla.models.backbones.siglip2_diffusion_backbone.SiglipVisionModel.from_pretrained", + return_value=fake_model, + ) as mock_from_pretrained: + backbone = SigLIP2DiffusionBackbone( + model_name="google/siglip2-base-patch16-256", + camera_names=_CAMERA_NAMES, + num_cameras=3, + per_view_output_dim=2, + freeze_backbone=True, + ) + + self.assertEqual(backbone.camera_names, _CAMERA_NAMES) + self.assertEqual(backbone.num_cameras, 3) + self.assertEqual(backbone.output_dim, 2) + self.assertEqual(backbone.joint_output_dim, 6) + self.assertIsNone(backbone.dataset_image_resize_shape) + self.assertEqual(backbone.eval_image_resize_shape, (256, 256)) + mock_from_pretrained.assert_called_once_with("google/siglip2-base-patch16-256") + self.assertTrue(all(not p.requires_grad for p in backbone.encoder.parameters())) + self.assertFalse(backbone.encoder.training) + + with torch.no_grad(): + backbone.view_projector.weight.zero_() + backbone.view_projector.bias.zero_() + backbone.view_projector.weight[0, 0] = 1.0 + backbone.view_projector.weight[1, 1] = 1.0 + + images = { + "r_vis": torch.full((1, 2, 3, 256, 256), 0.25), + "top": torch.full((1, 2, 3, 256, 256), 0.50), + "front": torch.full((1, 2, 3, 256, 256), 0.75), + } + output = backbone(images) + + self.assertEqual(output.shape, (1, 2, 6)) + self.assertEqual(len(fake_model.forward_calls), 3) + + expected_per_camera = [] + for cam_name in _CAMERA_NAMES: + img = images[cam_name].reshape(2, 3, 256, 256) + normalized = (img - 0.5) / 0.5 + expected_per_camera.append(normalized.mean(dim=(2, 3))[:, :2]) + expected = torch.cat(expected_per_camera, dim=-1).view(1, 2, 6) + self.assertTrue(torch.allclose(output, expected, atol=1e-6, rtol=1e-6)) + + for call, cam_name in zip(fake_model.forward_calls, _CAMERA_NAMES): + pixels = call["pixel_values"] + self.assertEqual(tuple(pixels.shape), (2, 3, 256, 256)) + self.assertTrue( + torch.allclose( + pixels, + (images[cam_name].reshape(2, 3, 256, 256) - 0.5) / 0.5, + ) + ) + + def test_forward_rejects_missing_required_camera(self): + from roboimi.vla.models.backbones.siglip2_diffusion_backbone import SigLIP2DiffusionBackbone + + backbone = SigLIP2DiffusionBackbone( + vision_model=_FakeSiglipVisionModel(hidden_size=4), + camera_names=_CAMERA_NAMES, + num_cameras=3, + ) + + with self.assertRaisesRegex(ValueError, "missing"): + backbone({ + "r_vis": torch.rand(1, 1, 3, 256, 256), + "top": torch.rand(1, 1, 3, 256, 256), + }) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_simple_robot_dataset_image_loading.py b/tests/test_simple_robot_dataset_image_loading.py new file mode 100644 index 0000000..b305275 --- /dev/null +++ b/tests/test_simple_robot_dataset_image_loading.py @@ -0,0 +1,81 @@ +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import h5py +import numpy as np + +from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset + + +class SimpleRobotDatasetImageLoadingTest(unittest.TestCase): + def _write_episode(self, dataset_dir: Path) -> None: + episode_path = dataset_dir / "episode_0.hdf5" + with h5py.File(episode_path, "w") as root: + root.create_dataset("action", data=np.arange(8, dtype=np.float32).reshape(4, 2)) + root.create_dataset( + "observations/qpos", + data=np.arange(16, dtype=np.float32).reshape(4, 4), + ) + root.create_dataset("task", data=np.array([b"sim_transfer"])) + root.create_dataset( + "observations/images/front", + data=np.arange(4 * 8 * 8 * 3, dtype=np.uint8).reshape(4, 8, 8, 3), + ) + + def test_getitem_only_resizes_observation_horizon_images(self): + with tempfile.TemporaryDirectory() as tmpdir: + dataset_dir = Path(tmpdir) + self._write_episode(dataset_dir) + dataset = SimpleRobotDataset( + dataset_dir, + obs_horizon=2, + pred_horizon=3, + camera_names=["front"], + ) + + resize_calls = [] + + def fake_resize(image, size, interpolation=None): + resize_calls.append( + { + "shape": tuple(image.shape), + "size": size, + "interpolation": interpolation, + } + ) + return image + + fake_cv2 = types.SimpleNamespace(INTER_LINEAR=1, resize=fake_resize) + + with mock.patch.dict(sys.modules, {"cv2": fake_cv2}): + sample = dataset[1] + + self.assertEqual(len(resize_calls), 2) + self.assertEqual(tuple(sample["observation.front"].shape), (2, 3, 8, 8)) + + def test_getitem_skips_resize_when_image_resize_shape_is_none(self): + with tempfile.TemporaryDirectory() as tmpdir: + dataset_dir = Path(tmpdir) + self._write_episode(dataset_dir) + dataset = SimpleRobotDataset( + dataset_dir, + obs_horizon=2, + pred_horizon=3, + camera_names=["front"], + image_resize_shape=None, + ) + + fake_cv2 = types.SimpleNamespace( + INTER_LINEAR=1, + resize=mock.Mock(side_effect=AssertionError("resize should be skipped when image_resize_shape=None")), + ) + + with mock.patch.dict(sys.modules, {"cv2": fake_cv2}): + sample = dataset[1] + + fake_cv2.resize.assert_not_called() + self.assertEqual(tuple(sample["observation.front"].shape), (2, 3, 8, 8)) diff --git a/tests/test_streaming_episode_writer.py b/tests/test_streaming_episode_writer.py new file mode 100644 index 0000000..0122d9d --- /dev/null +++ b/tests/test_streaming_episode_writer.py @@ -0,0 +1,79 @@ +import tempfile +import unittest +from pathlib import Path + +import h5py +import numpy as np + +from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter + + +class StreamingEpisodeWriterTest(unittest.TestCase): + def test_commit_persists_raw_action_and_resized_images(self): + camera_names = ["angle", "r_vis", "top", "front"] + raw_action_0 = np.arange(16, dtype=np.float32) + raw_action_1 = np.arange(16, dtype=np.float32) + 100.0 + qpos_0 = np.arange(16, dtype=np.float32) + 200.0 + qpos_1 = np.arange(16, dtype=np.float32) + 300.0 + + with tempfile.TemporaryDirectory() as tmpdir: + episode_path = Path(tmpdir) / "episode_0.hdf5" + writer = StreamingEpisodeWriter( + dataset_path=episode_path, + max_timesteps=2, + camera_names=camera_names, + image_size=(256, 256), + ) + + writer.append( + qpos=qpos_0, + action=raw_action_0, + images={ + cam: np.full((480, 640, 3), fill_value=idx + 1, dtype=np.uint8) + for idx, cam in enumerate(camera_names) + }, + ) + writer.append( + qpos=qpos_1, + action=raw_action_1, + images={ + cam: np.full((480, 640, 3), fill_value=idx + 11, dtype=np.uint8) + for idx, cam in enumerate(camera_names) + }, + ) + writer.commit() + + self.assertTrue(episode_path.exists()) + self.assertFalse(Path(str(episode_path) + ".tmp").exists()) + + with h5py.File(episode_path, "r") as root: + self.assertEqual(root["action"].shape, (2, 16)) + self.assertEqual(root["observations/qpos"].shape, (2, 16)) + np.testing.assert_allclose(root["action"][0], raw_action_0) + np.testing.assert_allclose(root["action"][1], raw_action_1) + np.testing.assert_allclose(root["observations/qpos"][0], qpos_0) + np.testing.assert_allclose(root["observations/qpos"][1], qpos_1) + for idx, cam_name in enumerate(camera_names): + dataset = root[f"observations/images/{cam_name}"] + self.assertEqual(dataset.shape, (2, 256, 256, 3)) + self.assertEqual(dataset.dtype, np.uint8) + self.assertTrue(np.all(dataset[0] == idx + 1)) + self.assertTrue(np.all(dataset[1] == idx + 11)) + + def test_discard_removes_temporary_file(self): + with tempfile.TemporaryDirectory() as tmpdir: + episode_path = Path(tmpdir) / "episode_0.hdf5" + writer = StreamingEpisodeWriter( + dataset_path=episode_path, + max_timesteps=1, + camera_names=["angle", "r_vis", "top", "front"], + image_size=(256, 256), + ) + writer.discard() + + self.assertFalse(episode_path.exists()) + self.assertFalse(Path(str(episode_path) + ".tmp").exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_train_vla_rollout_validation.py b/tests/test_train_vla_rollout_validation.py new file mode 100644 index 0000000..1dbdf9e --- /dev/null +++ b/tests/test_train_vla_rollout_validation.py @@ -0,0 +1,889 @@ +import os +import tempfile +import unittest +from copy import deepcopy +from pathlib import Path +from unittest import mock + +import numpy as np +import torch +from omegaconf import OmegaConf +from torch import nn + +from roboimi.demos.vla_scripts import eval_vla, train_vla + + +class _FakeDataset: + def __len__(self): + return 4 + + +class _FakeLoader: + def __init__(self, batch, length=1): + self._batches = [batch] * length + + def __len__(self): + return len(self._batches) + + def __iter__(self): + return iter(self._batches) + + +class _FakeOptimizer: + def __init__(self, lr=1e-3): + self.param_groups = [{'lr': lr}] + + def zero_grad(self): + return None + + def step(self): + return None + + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + del state_dict + return None + + +class _FakeScheduler: + def __init__(self): + self.step_calls = 0 + + def step(self): + self.step_calls += 1 + + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + del state_dict + return None + + +class _FakeProgressBar: + def __init__(self, iterable): + self._items = list(iterable) + self.postfix_calls = [] + + def __iter__(self): + return iter(self._items) + + def set_postfix(self, values): + self.postfix_calls.append(values) + + +class _FakeAgent(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.tensor(0.0)) + + def to(self, device): + del device + return self + + def compute_loss(self, agent_input): + del agent_input + return (self.weight - torch.tensor(0.5)).pow(2) + + def get_normalization_stats(self): + return {} + + +class _SequentialLossAgent(nn.Module): + def __init__(self, losses): + super().__init__() + self.weight = nn.Parameter(torch.tensor(0.0)) + self._losses = list(losses) + self._index = 0 + + def to(self, device): + del device + return self + + def compute_loss(self, agent_input): + del agent_input + loss_value = self._losses[self._index] + self._index += 1 + return (self.weight * 0) + torch.tensor(float(loss_value)) + + def get_normalization_stats(self): + return {} + + +class _FakeEvalAgent: + def __init__(self): + self.reset_calls = 0 + + def eval(self): + return self + + def to(self, device): + del device + return self + + def reset(self): + self.reset_calls += 1 + + def select_action(self, observation): + del observation + return torch.zeros(2) + + +class _FakeEvalEnv: + def reset(self, box_pos): + self.box_pos = box_pos + + def _get_image_obs(self): + return { + 'images': { + 'front': np.zeros((8, 8, 3), dtype=np.uint8), + } + } + + def _get_qpos_obs(self): + return {'qpos': np.zeros(4, dtype=np.float32)} + + def render(self): + raise AssertionError('render should not be called in this helper delegation test') + + +class TrainVLARolloutValidationTest(unittest.TestCase): + def test_default_train_config_uses_full_dataset_and_epoch_rollout_validation(self): + cfg = OmegaConf.load(Path('roboimi/vla/conf/config.yaml')) + + self.assertEqual(cfg.train.val_split, 0.0) + self.assertGreater(cfg.train.batch_size, 8) + self.assertGreater(float(cfg.train.lr), 5e-5) + self.assertGreater(cfg.train.num_workers, 8) + self.assertEqual(cfg.train.rollout_val_freq_epochs, 50) + + def test_training_passes_backbone_image_resize_override_to_dataset_instantiation(self): + cfg = OmegaConf.create( + { + 'agent': { + 'vision_backbone': { + 'dataset_image_resize_shape': None, + }, + 'normalization_type': 'min_max', + }, + 'data': { + 'dataset_dir': 'unused', + 'camera_names': ['front'], + }, + 'train': { + 'batch_size': 2, + 'lr': 1e-4, + 'max_steps': 0, + 'device': 'cpu', + 'disable_cudnn': False, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 42, + 'log_freq': 1, + 'save_freq': 10, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 0, + 'rollout_validate_on_checkpoint': False, + 'rollout_num_episodes': 1, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 1e-6, + 'weight_decay': 1e-5, + 'grad_clip': 1.0, + 'pretrained_ckpt': None, + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 1, + 'headless': True, + 'device': 'cpu', + 'verbose_action': False, + }, + 'experiment': {}, + } + ) + captured_dataset_kwargs = {} + + def fake_instantiate(config_node, **kwargs): + if config_node is cfg.data: + captured_dataset_kwargs.update(kwargs) + return _FakeDataset() + if config_node is cfg.agent: + return _FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(_dataset, *, shuffle, **_kwargs): + del shuffle, _kwargs + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=1, + ) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla, '_init_swanlab', return_value=None), \ + mock.patch.object(train_vla, '_finish_swanlab', return_value=None), \ + mock.patch.object(train_vla.torch, 'save', return_value=None): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertIn('image_resize_shape', captured_dataset_kwargs) + self.assertIsNone(captured_dataset_kwargs['image_resize_shape']) + + def test_eval_main_delegates_to_plain_run_eval_helper(self): + cfg = OmegaConf.create( + { + 'agent': {}, + 'eval': { + 'ckpt_path': 'checkpoints/vla_model_step_1.pt', + 'num_episodes': 1, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': True, + }, + } + ) + run_eval_mock = mock.Mock() + + with mock.patch.object(eval_vla, '_run_eval', run_eval_mock, create=True), \ + mock.patch.object(eval_vla, 'load_checkpoint', return_value=(_FakeEvalAgent(), None)), \ + mock.patch.object(eval_vla, 'make_sim_env', return_value=_FakeEvalEnv()), \ + mock.patch.object(eval_vla, 'sample_transfer_pose', return_value=np.zeros(3)), \ + mock.patch.object(eval_vla, 'execute_policy_action'), \ + mock.patch.object(eval_vla, 'tqdm', side_effect=lambda iterable, **kwargs: iterable): + eval_vla.main.__wrapped__(cfg) + + run_eval_mock.assert_called_once_with(cfg) + + def test_run_training_rollout_validation_runs_every_50_epochs_and_uses_avg_reward_metric(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 1, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 100, + 'log_freq': 1, + 'save_freq': 1000, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 50, + 'rollout_num_episodes': 3, + }, + 'data': { + 'camera_names': ['front'], + }, + 'agent': { + '_target_': 'fake.agent', + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 99, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': False, + }, + } + ) + agent = _FakeAgent() + rollout_mock = mock.Mock( + side_effect=[ + { + 'avg_reward': 2.0, + 'episodes': [ + { + 'episode_index': 0, + 'artifact_paths': {'trajectory_image': 'artifacts/epoch_49_front.png'}, + }, + ], + }, + { + 'avg_reward': 1.0, + 'episodes': [ + { + 'episode_index': 0, + 'artifact_paths': {'trajectory_image': 'artifacts/epoch_99_front.png'}, + }, + ], + }, + ] + ) + swanlab_log_mock = mock.Mock() + saved_checkpoints = [] + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(_dataset, *, shuffle, **_kwargs): + del shuffle, _kwargs + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=1, + ) + + def fake_torch_save(payload, path): + saved_checkpoints.append((str(path), deepcopy(payload))) + return None + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla, '_log_to_swanlab', swanlab_log_mock), \ + mock.patch.object(train_vla.torch, 'save', side_effect=fake_torch_save), \ + mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True), \ + mock.patch.object(eval_vla.main, '__wrapped__', side_effect=AssertionError('training hook should call eval_vla._run_eval')): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual(rollout_mock.call_count, 2) + first_rollout_cfg = rollout_mock.call_args_list[0].args[0] + second_rollout_cfg = rollout_mock.call_args_list[1].args[0] + self.assertTrue(first_rollout_cfg.eval.ckpt_path.endswith('checkpoints/vla_model_step_49.pt')) + self.assertTrue(second_rollout_cfg.eval.ckpt_path.endswith('checkpoints/vla_model_step_99.pt')) + self.assertEqual(first_rollout_cfg.eval.num_episodes, 3) + self.assertTrue(first_rollout_cfg.eval.headless) + self.assertEqual(first_rollout_cfg.eval.device, 'cpu') + self.assertFalse(first_rollout_cfg.eval.verbose_action) + self.assertFalse(first_rollout_cfg.eval.record_video) + self.assertTrue(first_rollout_cfg.eval.save_trajectory_image) + self.assertEqual(first_rollout_cfg.eval.trajectory_image_camera_name, 'front') + self.assertEqual(cfg.eval.ckpt_path, 'unused.pt') + self.assertEqual(cfg.eval.num_episodes, 99) + self.assertFalse(cfg.eval.headless) + self.assertEqual(cfg.eval.device, 'cpu') + self.assertFalse(cfg.eval.verbose_action) + self.assertNotIn('save_trajectory_image', cfg.eval) + self.assertNotIn('trajectory_image_camera_name', cfg.eval) + + rollout_reward_logs = [ + call.args[1]['rollout/avg_reward'] + for call in swanlab_log_mock.call_args_list + if len(call.args) >= 2 and 'rollout/avg_reward' in call.args[1] + ] + self.assertEqual(rollout_reward_logs, [2.0, 1.0]) + + best_model_saves = [ + payload for path, payload in saved_checkpoints + if path.endswith('checkpoints/vla_model_best.pt') + ] + self.assertEqual(len(best_model_saves), 1) + self.assertEqual(best_model_saves[0]['rollout_avg_reward'], 2.0) + + def test_run_training_keeps_loss_based_best_checkpoint_until_first_rollout_metric_exists(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 1, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 5, + 'log_freq': 1, + 'save_freq': 2, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 50, + 'rollout_num_episodes': 3, + }, + 'data': { + 'camera_names': ['front'], + }, + 'agent': { + '_target_': 'fake.agent', + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 99, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': False, + }, + } + ) + saved_checkpoints = [] + rollout_mock = mock.Mock() + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _FakeDataset() + if config_node is cfg.agent: + return _FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(_dataset, *, shuffle, **_kwargs): + del shuffle, _kwargs + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=5, + ) + + def fake_torch_save(payload, path): + saved_checkpoints.append((str(path), deepcopy(payload))) + return None + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla.torch, 'save', side_effect=fake_torch_save), \ + mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual(rollout_mock.call_count, 0) + best_model_saves = [ + payload for path, payload in saved_checkpoints + if path.endswith('checkpoints/vla_model_best.pt') + ] + self.assertEqual(len(best_model_saves), 1) + self.assertIsNone(best_model_saves[0]['rollout_avg_reward']) + + def test_run_training_disables_drop_last_when_train_set_is_smaller_than_batch_size(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 8, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 1, + 'log_freq': 1, + 'save_freq': 10, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 50, + 'rollout_num_episodes': 3, + }, + 'data': { + 'camera_names': ['front'], + }, + 'agent': { + '_target_': 'fake.agent', + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 99, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': False, + }, + } + ) + dataloader_calls = [] + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _FakeDataset() + if config_node is cfg.agent: + return _FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(dataset, *, shuffle, drop_last, **_kwargs): + dataloader_calls.append({ + 'shuffle': shuffle, + 'drop_last': drop_last, + 'dataset_len': len(dataset), + }) + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=1, + ) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla.torch, 'save', return_value=None): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + train_loader_calls = [call for call in dataloader_calls if call['shuffle']] + self.assertEqual(len(train_loader_calls), 1) + self.assertFalse(train_loader_calls[0]['drop_last']) + + def test_run_training_disables_persistent_workers_for_train_and_val_loaders(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 2, + 'num_workers': 2, + 'val_split': 0.25, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 1, + 'log_freq': 1, + 'save_freq': 10, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 50, + 'rollout_num_episodes': 3, + }, + 'data': { + 'camera_names': ['front'], + }, + 'agent': { + '_target_': 'fake.agent', + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 99, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': False, + }, + } + ) + dataloader_calls = [] + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _FakeDataset() + if config_node is cfg.agent: + return _FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(_dataset, *, shuffle, persistent_workers, num_workers, **_kwargs): + dataloader_calls.append({ + 'shuffle': shuffle, + 'num_workers': num_workers, + 'persistent_workers': persistent_workers, + }) + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=1, + ) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla.torch, 'save', return_value=None): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual(len(dataloader_calls), 2) + self.assertEqual([call['shuffle'] for call in dataloader_calls], [True, False]) + self.assertTrue(all(call['num_workers'] == 2 for call in dataloader_calls)) + self.assertTrue(all(call['persistent_workers'] is False for call in dataloader_calls)) + + def test_run_training_uses_loss_best_until_first_rollout_then_prefers_rollout_reward(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 1, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 6, + 'log_freq': 1, + 'save_freq': 1, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 2, + 'rollout_num_episodes': 1, + }, + 'data': { + 'camera_names': ['front'], + }, + 'agent': { + '_target_': 'fake.agent', + }, + 'eval': { + 'ckpt_path': 'unused.pt', + 'num_episodes': 99, + 'max_timesteps': 1, + 'device': 'cpu', + 'task_name': 'sim_transfer', + 'camera_names': ['front'], + 'use_smoothing': False, + 'smooth_alpha': 0.3, + 'verbose_action': False, + 'headless': False, + }, + } + ) + agent = _SequentialLossAgent([10, 9, 8, 7, 6, 5]) + rollout_mock = mock.Mock(return_value={'avg_reward': 1.0}) + saved_checkpoints = [] + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(_dataset, *, shuffle, **_kwargs): + del _kwargs + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=2 if shuffle else 1, + ) + + def fake_torch_save(payload, path): + saved_checkpoints.append((str(path), deepcopy(payload))) + return None + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla.torch, 'save', side_effect=fake_torch_save), \ + mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + best_model_saves = [ + (payload['step'], payload['rollout_avg_reward']) + for path, payload in saved_checkpoints + if path.endswith('checkpoints/vla_model_best.pt') + ] + self.assertEqual( + best_model_saves, + [ + (1, None), + (2, None), + (3, None), + (3, 1.0), + ], + ) + self.assertEqual(rollout_mock.call_count, 1) + + def test_run_training_keeps_tiny_train_dataset_batch_when_batch_size_is_larger(self): + cfg = OmegaConf.create( + { + 'train': { + 'device': 'cpu', + 'batch_size': 8, + 'num_workers': 0, + 'val_split': 0.0, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 1, + 'log_freq': 1, + 'save_freq': 1000, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': False, + 'rollout_val_freq_epochs': 0, + }, + 'data': { + 'camera_names': ['front'], + }, + 'agent': { + '_target_': 'fake.agent', + }, + } + ) + agent = _FakeAgent() + dataloader_calls = [] + saved_checkpoints = [] + + class _TinyDataset: + def __len__(self): + return 1 + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return _TinyDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_dataloader(dataset, *, drop_last, shuffle, **_kwargs): + del _kwargs + dataloader_calls.append( + { + 'shuffle': shuffle, + 'drop_last': drop_last, + 'dataset_len': len(dataset), + } + ) + loader_length = 0 if drop_last and len(dataset) < cfg.train.batch_size else 1 + return _FakeLoader( + { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + }, + length=loader_length, + ) + + def fake_torch_save(payload, path): + saved_checkpoints.append((str(path), deepcopy(payload))) + return None + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \ + mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \ + mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \ + mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \ + mock.patch.object(train_vla.torch, 'save', side_effect=fake_torch_save): + train_vla._run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual( + dataloader_calls[0], + { + 'shuffle': True, + 'drop_last': False, + 'dataset_len': 1, + }, + ) + self.assertEqual(len(saved_checkpoints), 1) + self.assertTrue(saved_checkpoints[0][0].endswith('checkpoints/vla_model_final.pt')) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_train_vla_swanlab_logging.py b/tests/test_train_vla_swanlab_logging.py new file mode 100644 index 0000000..3918e9b --- /dev/null +++ b/tests/test_train_vla_swanlab_logging.py @@ -0,0 +1,826 @@ +import importlib +import importlib.util +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import torch +from torch import nn + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_TRAIN_VLA_PATH = _REPO_ROOT / 'roboimi/demos/vla_scripts/train_vla.py' +_CONFIG_PATH = _REPO_ROOT / 'roboimi/vla/conf/config.yaml' + + +class AttrDict(dict): + def __getattr__(self, name): + try: + return self[name] + except KeyError as exc: + raise AttributeError(name) from exc + + def __setattr__(self, name, value): + self[name] = value + + +def _to_attrdict(value): + if isinstance(value, dict): + return AttrDict({key: _to_attrdict(item) for key, item in value.items()}) + if isinstance(value, list): + return [_to_attrdict(item) for item in value] + return value + + +class FakeDataset: + def __len__(self): + return 4 + + +class FakeLoader: + def __init__(self, batch): + self.batch = batch + + def __len__(self): + return 1 + + def __iter__(self): + return iter((self.batch,)) + + +class FakeScheduler: + def __init__(self): + self.step_calls = 0 + + def step(self): + self.step_calls += 1 + + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + return None + + +class FakeOptimizer: + def __init__(self, lr=1e-3): + self.param_groups = [{'lr': lr}] + self.loaded_state_dict = None + + def zero_grad(self): + return None + + def step(self): + return None + + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + self.loaded_state_dict = state_dict + return None + + +class FakeProgressBar: + def __init__(self, iterable): + self._items = list(iterable) + self.postfix_calls = [] + + def __iter__(self): + return iter(self._items) + + def set_postfix(self, values): + self.postfix_calls.append(values) + + +class FakeAgent(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.tensor(0.0)) + + def to(self, device): + return self + + def compute_loss(self, agent_input): + del agent_input + target = torch.tensor(0.25 if self.training else 0.1) + return (self.weight - target).pow(2) + + def get_normalization_stats(self): + return {} + + +class FakeSwanLab: + def __init__(self, init_error=None, log_errors=None, finish_error=None, image_errors=None): + self.init_error = init_error + self.log_errors = list(log_errors or []) + self.finish_error = finish_error + self.image_errors = list(image_errors or []) + self.init_calls = [] + self.log_calls = [] + self.finish_calls = 0 + self.image_calls = [] + + def init(self, project, experiment_name=None, config=None): + self.init_calls.append({ + 'project': project, + 'experiment_name': experiment_name, + 'config': config, + }) + if self.init_error is not None: + raise self.init_error + return object() + + def log(self, payload, step=None): + self.log_calls.append((dict(payload), step)) + if self.log_errors: + raise self.log_errors.pop(0) + + def Image(self, path, caption=None): + self.image_calls.append({ + 'path': path, + 'caption': caption, + }) + if self.image_errors: + raise self.image_errors.pop(0) + return { + 'path': path, + 'caption': caption, + } + + def finish(self): + self.finish_calls += 1 + if self.finish_error is not None: + raise self.finish_error + + +class TrainVLASwanLabLoggingTest(unittest.TestCase): + def test_default_config_keeps_swanlab_opt_in(self): + config_text = _CONFIG_PATH.read_text(encoding='utf-8') + self.assertIn('use_swanlab: false', config_text) + + def test_log_rollout_trajectory_images_to_swanlab_uploads_episode_artifacts(self): + module = self._load_train_vla_module() + fake_swanlab = FakeSwanLab() + + module._log_rollout_trajectory_images_to_swanlab( + fake_swanlab, + { + 'episodes': [ + { + 'episode_index': 0, + 'artifact_paths': {'trajectory_image': 'artifacts/episode_0_front.png'}, + }, + { + 'episode_index': 3, + 'artifact_paths': {'trajectory_image': 'artifacts/episode_3_front.png'}, + }, + { + 'episode_index': 7, + 'artifact_paths': {'trajectory_image': None}, + }, + { + 'episode_index': 8, + 'artifact_paths': {}, + }, + ], + }, + step=12, + context_label='epoch 1 rollout', + ) + + self.assertEqual( + fake_swanlab.image_calls, + [ + { + 'path': 'artifacts/episode_0_front.png', + 'caption': 'epoch 1 rollout trajectory image - episode 0 (front)', + }, + { + 'path': 'artifacts/episode_3_front.png', + 'caption': 'epoch 1 rollout trajectory image - episode 3 (front)', + }, + ], + ) + self.assertIn( + ( + { + 'rollout/trajectory_image_episode_0': { + 'path': 'artifacts/episode_0_front.png', + 'caption': 'epoch 1 rollout trajectory image - episode 0 (front)', + }, + 'rollout/trajectory_image_episode_3': { + 'path': 'artifacts/episode_3_front.png', + 'caption': 'epoch 1 rollout trajectory image - episode 3 (front)', + }, + }, + 12, + ), + fake_swanlab.log_calls, + ) + + def test_log_rollout_trajectory_images_to_swanlab_is_best_effort(self): + module = self._load_train_vla_module() + fake_swanlab = FakeSwanLab(image_errors=[RuntimeError('decode failed')]) + + with mock.patch.object(module.log, 'warning') as warning_mock: + module._log_rollout_trajectory_images_to_swanlab( + fake_swanlab, + { + 'episodes': [ + { + 'episode_index': 0, + 'artifact_paths': {'trajectory_image': 'artifacts/bad_episode.png'}, + }, + { + 'episode_index': 1, + 'artifact_paths': {'trajectory_image': 'artifacts/good_episode.png'}, + }, + ], + }, + step=7, + context_label='checkpoint rollout', + ) + + self.assertEqual( + fake_swanlab.image_calls, + [ + { + 'path': 'artifacts/bad_episode.png', + 'caption': 'checkpoint rollout trajectory image - episode 0 (front)', + }, + { + 'path': 'artifacts/good_episode.png', + 'caption': 'checkpoint rollout trajectory image - episode 1 (front)', + }, + ], + ) + self.assertIn( + ( + { + 'rollout/trajectory_image_episode_1': { + 'path': 'artifacts/good_episode.png', + 'caption': 'checkpoint rollout trajectory image - episode 1 (front)', + }, + }, + 7, + ), + fake_swanlab.log_calls, + ) + warning_messages = [call.args[0] for call in warning_mock.call_args_list] + self.assertTrue( + any('SwanLab rollout trajectory image upload prep failed' in message for message in warning_messages) + ) + + def _load_train_vla_module(self): + hydra_module = types.ModuleType('hydra') + hydra_utils_module = types.ModuleType('hydra.utils') + hydra_utils_module.instantiate = lambda *args, **kwargs: None + + def hydra_main(**_kwargs): + def decorator(func): + return func + return decorator + + hydra_module.main = hydra_main + hydra_module.utils = hydra_utils_module + + class OmegaConfStub: + _resolvers = {} + + @classmethod + def has_resolver(cls, name): + return name in cls._resolvers + + @classmethod + def register_new_resolver(cls, name, resolver): + cls._resolvers[name] = resolver + + @staticmethod + def to_yaml(_cfg): + return 'stub-config' + + @staticmethod + def to_container(cfg, resolve=False): + del resolve + return dict(cfg) + + @staticmethod + def create(cfg): + return _to_attrdict(cfg) + + omegaconf_module = types.ModuleType('omegaconf') + omegaconf_module.DictConfig = dict + omegaconf_module.OmegaConf = OmegaConfStub + + module_name = 'train_vla_swanlab_test_module' + spec = importlib.util.spec_from_file_location(module_name, _TRAIN_VLA_PATH) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict( + sys.modules, + { + 'hydra': hydra_module, + 'hydra.utils': hydra_utils_module, + 'omegaconf': omegaconf_module, + }, + ): + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + def _make_cfg(self, *, use_swanlab=True, swanlab_run_name='smoke-run'): + return AttrDict( + train=AttrDict( + device='cpu', + batch_size=2, + num_workers=0, + val_split=0.25, + seed=0, + lr=1e-3, + max_steps=2, + log_freq=1, + save_freq=1, + warmup_steps=1, + scheduler_type='constant', + min_lr=0.0, + grad_clip=1.0, + weight_decay=0.0, + pretrained_ckpt=None, + resume_ckpt=None, + use_swanlab=use_swanlab, + swanlab_project='roboimi-vla-tests', + swanlab_run_name=swanlab_run_name, + ), + data=AttrDict( + camera_names=('front',), + ), + agent=AttrDict( + _target_='fake.agent', + ), + eval=AttrDict( + ckpt_path='unused.pt', + num_episodes=1, + max_timesteps=1, + device='cpu', + task_name='sim_transfer', + camera_names=('front',), + use_smoothing=False, + smooth_alpha=0.3, + verbose_action=False, + headless=False, + ), + ) + + def _get_run_training(self, module): + run_training = getattr(module, '_run_training', None) + self.assertIsNotNone(run_training, 'Expected train_vla.py to expose a _run_training(cfg) helper') + return run_training + + def _make_batch(self): + return { + 'observation.front': torch.zeros(1, 3, 2, 2), + 'observation.state': torch.zeros(1, 4), + 'action': torch.zeros(1, 2), + 'action_is_pad': torch.zeros(1, 1, dtype=torch.bool), + } + + def _loader_factory(self): + train_batch = self._make_batch() + val_batch = self._make_batch() + + def factory(_dataset, *, shuffle, **_kwargs): + return FakeLoader(train_batch if shuffle else val_batch) + + return factory + + def test_run_training_logs_metrics_and_checkpoint_paths_to_swanlab(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + agent = FakeAgent() + fake_swanlab = FakeSwanLab() + real_import_module = importlib.import_module + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module): + run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual( + fake_swanlab.init_calls, + [{ + 'project': 'roboimi-vla-tests', + 'experiment_name': 'smoke-run', + 'config': { + 'train': { + 'device': 'cpu', + 'batch_size': 2, + 'num_workers': 0, + 'val_split': 0.25, + 'seed': 0, + 'lr': 1e-3, + 'max_steps': 2, + 'log_freq': 1, + 'save_freq': 1, + 'warmup_steps': 1, + 'scheduler_type': 'constant', + 'min_lr': 0.0, + 'grad_clip': 1.0, + 'weight_decay': 0.0, + 'pretrained_ckpt': None, + 'resume_ckpt': None, + 'use_swanlab': True, + 'swanlab_project': 'roboimi-vla-tests', + 'swanlab_run_name': 'smoke-run', + }, + 'data': { + 'camera_names': ('front',), + }, + 'agent': { + '_target_': 'fake.agent', + }, + }, + }], + ) + + logged_keys = set().union(*(payload.keys() for payload, _step in fake_swanlab.log_calls)) + self.assertTrue( + { + 'train/loss', + 'train/lr', + 'train/best_loss', + 'train/step', + 'val/loss', + 'final/checkpoint_path', + 'final/best_checkpoint_path', + }.issubset(logged_keys) + ) + + final_payload, final_step = fake_swanlab.log_calls[-1] + self.assertEqual(final_step, cfg.train.max_steps) + self.assertTrue(final_payload['final/checkpoint_path'].endswith('checkpoints/vla_model_final.pt')) + self.assertTrue(final_payload['final/best_checkpoint_path'].endswith('checkpoints/vla_model_best.pt')) + self.assertEqual(fake_swanlab.finish_calls, 1) + + def test_run_training_skips_swanlab_when_disabled(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg(use_swanlab=False) + agent = FakeAgent() + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module.importlib, 'import_module', side_effect=AssertionError('swanlab import should not run')): + run_training(cfg) + finally: + os.chdir(previous_cwd) + + def test_run_training_finishes_swanlab_when_exception_happens_after_init(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + fake_swanlab = FakeSwanLab() + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(module, 'instantiate', side_effect=RuntimeError('dataset boom')), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module): + with self.assertRaisesRegex(RuntimeError, 'dataset boom'): + run_training(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual(fake_swanlab.finish_calls, 1) + + def test_run_training_warns_and_continues_when_swanlab_log_and_finish_fail(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + agent = FakeAgent() + fake_swanlab = FakeSwanLab( + log_errors=[RuntimeError('log backend hiccup')], + finish_error=RuntimeError('finish backend hiccup'), + ) + real_import_module = importlib.import_module + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module), \ + mock.patch.object(module.log, 'warning') as warning_mock: + run_training(cfg) + finally: + os.chdir(previous_cwd) + + warning_messages = [call.args[0] for call in warning_mock.call_args_list] + self.assertTrue(any('SwanLab log failed' in message for message in warning_messages)) + self.assertTrue(any('SwanLab finish failed' in message for message in warning_messages)) + self.assertEqual(fake_swanlab.finish_calls, 1) + + def test_run_training_resume_restores_best_rollout_baseline_from_best_checkpoint(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + cfg.train.max_steps = 2 + cfg.train.save_freq = 1 + cfg.train.rollout_validate_on_checkpoint = True + fake_swanlab = FakeSwanLab() + fake_optimizer = FakeOptimizer(lr=cfg.train.lr) + fake_scheduler = FakeScheduler() + real_import_module = importlib.import_module + saved_paths = [] + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + checkpoint_dir = Path('checkpoints') + checkpoint_dir.mkdir() + resume_path = checkpoint_dir / 'vla_model_step_0.pt' + resume_path.write_bytes(b'resume') + best_path = checkpoint_dir / 'vla_model_best.pt' + best_path.write_bytes(b'best') + cfg.train.resume_ckpt = str(resume_path) + + resume_checkpoint_state = { + 'step': 0, + 'model_state_dict': FakeAgent().state_dict(), + 'optimizer_state_dict': {}, + 'scheduler_state_dict': {}, + 'loss': 0.5, + 'val_loss': 0.25, + } + best_checkpoint_state = { + 'step': 0, + 'model_state_dict': FakeAgent().state_dict(), + 'optimizer_state_dict': {}, + 'scheduler_state_dict': {}, + 'loss': 0.5, + 'val_loss': 0.25, + 'rollout_avg_reward': 5.0, + } + + def fake_torch_load(path, map_location=None): + del map_location + path = Path(path).resolve() + if path == resume_path.resolve(): + return resume_checkpoint_state + if path == best_path.resolve(): + return best_checkpoint_state + raise AssertionError(f'unexpected load path: {path}') + + def fake_torch_save(payload, path): + saved_paths.append(str(path)) + return None + + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \ + mock.patch.object(module, 'build_training_optimizer', return_value=fake_optimizer), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=fake_scheduler), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \ + mock.patch.object(module.torch, 'save', side_effect=fake_torch_save), \ + mock.patch.object(module.torch, 'load', side_effect=fake_torch_load), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module), \ + mock.patch('roboimi.demos.vla_scripts.eval_vla._run_eval', return_value={'avg_reward': 3.0}): + run_training(cfg) + finally: + os.chdir(previous_cwd) + + final_payload, final_step = fake_swanlab.log_calls[-1] + self.assertEqual(final_step, cfg.train.max_steps) + self.assertTrue(final_payload['final/best_checkpoint_path'].endswith('checkpoints/vla_model_best.pt')) + self.assertFalse(any(path.endswith('checkpoints/vla_model_best.pt') for path in saved_paths)) + + def test_run_training_resume_ignores_best_checkpoint_without_rollout_metric(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + cfg.train.max_steps = 1 + fake_swanlab = FakeSwanLab() + fake_optimizer = FakeOptimizer(lr=cfg.train.lr) + fake_scheduler = FakeScheduler() + real_import_module = importlib.import_module + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + checkpoint_dir = Path('checkpoints') + checkpoint_dir.mkdir() + resume_path = checkpoint_dir / 'vla_model_step_0.pt' + resume_path.write_bytes(b'resume') + best_path = checkpoint_dir / 'vla_model_best.pt' + best_path.write_bytes(b'stale') + cfg.train.resume_ckpt = str(resume_path) + + resume_checkpoint_state = { + 'step': 0, + 'model_state_dict': FakeAgent().state_dict(), + 'optimizer_state_dict': {}, + 'scheduler_state_dict': {}, + 'loss': 0.5, + 'val_loss': 0.25, + } + stale_best_checkpoint_state = { + 'step': 0, + 'model_state_dict': FakeAgent().state_dict(), + 'optimizer_state_dict': {}, + 'scheduler_state_dict': {}, + 'loss': 0.4, + 'val_loss': 0.2, + } + + def fake_torch_load(path, map_location=None): + del map_location + path = Path(path).resolve() + if path == resume_path.resolve(): + return resume_checkpoint_state + if path == best_path.resolve(): + return stale_best_checkpoint_state + raise AssertionError(f'unexpected load path: {path}') + + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \ + mock.patch.object(module, 'build_training_optimizer', return_value=fake_optimizer), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=fake_scheduler), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module.torch, 'load', side_effect=fake_torch_load), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module): + run_training(cfg) + finally: + os.chdir(previous_cwd) + + final_payload, final_step = fake_swanlab.log_calls[-1] + self.assertEqual(final_step, cfg.train.max_steps) + self.assertEqual(final_payload['final/best_checkpoint_path'], 'checkpoints/vla_model_step_0.pt') + + def test_run_training_ignores_stale_best_checkpoint_file_on_fresh_non_resume_run(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + cfg.train.max_steps = 1 + fake_swanlab = FakeSwanLab() + real_import_module = importlib.import_module + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return FakeAgent() + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + checkpoint_dir = Path('checkpoints') + checkpoint_dir.mkdir() + (checkpoint_dir / 'vla_model_best.pt').write_bytes(b'stale-best') + + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module): + run_training(cfg) + finally: + os.chdir(previous_cwd) + + final_payload, final_step = fake_swanlab.log_calls[-1] + self.assertEqual(final_step, cfg.train.max_steps) + self.assertEqual(final_payload['final/best_checkpoint_path'], '') + + def test_run_training_fails_fast_when_swanlab_import_is_unavailable(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == 'swanlab': + raise ImportError('missing swanlab') + return real_import_module(name, package) + + with mock.patch.object(module, 'instantiate', side_effect=AssertionError('instantiate should not run')), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module): + with self.assertRaisesRegex(RuntimeError, 'SwanLab'): + run_training(cfg) + + def test_run_training_fails_fast_when_swanlab_init_fails(self): + module = self._load_train_vla_module() + run_training = self._get_run_training(module) + cfg = self._make_cfg() + fake_swanlab = FakeSwanLab(init_error=RuntimeError('not logged in')) + real_import_module = importlib.import_module + + def fake_import_module(name, package=None): + if name == 'swanlab': + return fake_swanlab + return real_import_module(name, package) + + with mock.patch.object(module, 'instantiate', side_effect=AssertionError('instantiate should not run')), \ + mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module): + with self.assertRaisesRegex(RuntimeError, 'not logged in'): + run_training(cfg) + + self.assertEqual(fake_swanlab.finish_calls, 0) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_train_vla_transformer_optimizer.py b/tests/test_train_vla_transformer_optimizer.py new file mode 100644 index 0000000..56ee107 --- /dev/null +++ b/tests/test_train_vla_transformer_optimizer.py @@ -0,0 +1,424 @@ +import importlib.util +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import torch +from torch import nn + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_TRAIN_VLA_PATH = _REPO_ROOT / 'roboimi/demos/vla_scripts/train_vla.py' + + +class AttrDict(dict): + def __getattr__(self, name): + try: + return self[name] + except KeyError as exc: + raise AttributeError(name) from exc + + def __setattr__(self, name, value): + self[name] = value + + +class FakeDataset: + def __len__(self): + return 4 + + +class FakeLoader: + def __len__(self): + return 1 + + def __iter__(self): + return iter(()) + + +class FakeScheduler: + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + return None + + +class RecordingAdamW: + created = [] + + def __init__(self, params, lr, weight_decay): + self.lr = lr + self.weight_decay = weight_decay + self.param_groups = self._normalize_param_groups(params, lr, weight_decay) + RecordingAdamW.created.append(self) + + @staticmethod + def _normalize_param_groups(params, lr, weight_decay): + if isinstance(params, (list, tuple)) and params and isinstance(params[0], dict): + groups = [] + for group in params: + normalized = dict(group) + normalized['params'] = list(group['params']) + normalized.setdefault('lr', lr) + groups.append(normalized) + return groups + + return [{ + 'params': list(params), + 'lr': lr, + 'weight_decay': weight_decay, + }] + + def state_dict(self): + return {} + + def load_state_dict(self, state_dict): + return None + + +class RecordingTransformerHead(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + self.norm = nn.LayerNorm(4) + self.optim_group_calls = [] + + def get_optim_groups(self, weight_decay): + self.optim_group_calls.append(weight_decay) + return [ + { + 'params': [self.proj.weight], + 'weight_decay': weight_decay, + }, + { + 'params': [self.proj.bias, self.norm.weight, self.norm.bias], + 'weight_decay': 0.0, + }, + ] + + +class FakeIMFAgent(nn.Module): + def __init__(self): + super().__init__() + self.head_type = 'imf_transformer' + self.noise_pred_net = RecordingTransformerHead() + self.backbone = nn.Linear(4, 3) + self.adapter = nn.Linear(3, 2, bias=False) + + +class FakeTransformerAgent(nn.Module): + def __init__(self, *, head_type='transformer'): + super().__init__() + self.head_type = head_type + self.noise_pred_net = RecordingTransformerHead() + self.backbone = nn.Linear(4, 3) + self.adapter = nn.Linear(3, 2, bias=False) + self.frozen = nn.Linear(2, 2) + for param in self.frozen.parameters(): + param.requires_grad = False + + def to(self, device): + return self + + def get_normalization_stats(self): + return {} + + +class TrainVLATransformerOptimizerTest(unittest.TestCase): + def setUp(self): + RecordingAdamW.created = [] + + def _load_train_vla_module(self): + hydra_module = types.ModuleType('hydra') + hydra_utils_module = types.ModuleType('hydra.utils') + hydra_utils_module.instantiate = lambda *args, **kwargs: None + + def hydra_main(**_kwargs): + def decorator(func): + return func + return decorator + + hydra_module.main = hydra_main + hydra_module.utils = hydra_utils_module + + class OmegaConfStub: + _resolvers = {} + + @classmethod + def has_resolver(cls, name): + return name in cls._resolvers + + @classmethod + def register_new_resolver(cls, name, resolver): + cls._resolvers[name] = resolver + + @staticmethod + def to_yaml(_cfg): + return 'stub-config' + + omegaconf_module = types.ModuleType('omegaconf') + omegaconf_module.DictConfig = dict + omegaconf_module.OmegaConf = OmegaConfStub + + module_name = 'train_vla_optimizer_test_module' + spec = importlib.util.spec_from_file_location(module_name, _TRAIN_VLA_PATH) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict( + sys.modules, + { + 'hydra': hydra_module, + 'hydra.utils': hydra_utils_module, + 'omegaconf': omegaconf_module, + }, + ): + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + def _make_cfg(self): + return AttrDict( + train=AttrDict( + device='cpu', + batch_size=2, + num_workers=0, + val_split=0, + seed=0, + lr=1e-4, + max_steps=0, + log_freq=1, + save_freq=100, + warmup_steps=1, + scheduler_type='constant', + min_lr=0.0, + grad_clip=1.0, + weight_decay=0.123, + pretrained_ckpt=None, + resume_ckpt=None, + ), + data=AttrDict( + camera_names=('front',), + ), + agent=AttrDict( + _target_='fake.agent', + ), + ) + + def _group_names(self, agent, optimizer): + names_by_param_id = {id(param): name for name, param in agent.named_parameters()} + return [ + {names_by_param_id[id(param)] for param in group['params']} + for group in optimizer.param_groups + ] + + def test_clean_ld_preload_value_removes_problematic_nxegl_entry(self): + module = self._load_train_vla_module() + + cleaned, changed = module._clean_ld_preload_value( + '/usr/lib/libfoo.so /usr/NX/lib/libnxegl.so /usr/lib/libbar.so' + ) + + self.assertTrue(changed) + self.assertEqual(cleaned, '/usr/lib/libfoo.so /usr/lib/libbar.so') + + def test_clean_ld_preload_value_leaves_safe_entries_unchanged(self): + module = self._load_train_vla_module() + + cleaned, changed = module._clean_ld_preload_value('/usr/lib/libfoo.so /usr/lib/libbar.so') + + self.assertFalse(changed) + self.assertEqual(cleaned, '/usr/lib/libfoo.so /usr/lib/libbar.so') + + + def test_configure_cuda_runtime_can_disable_cudnn_for_training(self): + module = self._load_train_vla_module() + cfg = AttrDict(train=AttrDict(device='cuda', disable_cudnn=True)) + + original = module.torch.backends.cudnn.enabled + try: + module.torch.backends.cudnn.enabled = True + module._configure_cuda_runtime(cfg) + self.assertFalse(module.torch.backends.cudnn.enabled) + finally: + module.torch.backends.cudnn.enabled = original + + + def test_resolve_run_output_dir_prefers_hydra_runtime_output_dir(self): + module = self._load_train_vla_module() + hydra_core_module = types.ModuleType('hydra.core') + hydra_hydra_config_module = types.ModuleType('hydra.core.hydra_config') + + class _Runtime: + output_dir = '/tmp/hydra-output' + + class _Cfg: + runtime = _Runtime() + + class HydraConfigStub: + @staticmethod + def initialized(): + return True + @staticmethod + def get(): + return _Cfg() + + hydra_hydra_config_module.HydraConfig = HydraConfigStub + with mock.patch.dict(sys.modules, { + 'hydra.core': hydra_core_module, + 'hydra.core.hydra_config': hydra_hydra_config_module, + }): + output_dir = module._resolve_run_output_dir() + + self.assertEqual(Path(output_dir).resolve(), Path('/tmp/hydra-output').resolve()) + + + def test_train_script_uses_file_based_repo_root_on_sys_path(self): + module = self._load_train_vla_module() + + fake_sys_path = ['/tmp/site-packages', '/another/path'] + with mock.patch.object(module.sys, 'path', fake_sys_path): + repo_root = module._ensure_repo_root_on_syspath() + + self.assertEqual(Path(repo_root).resolve(), _REPO_ROOT.resolve()) + self.assertEqual(Path(fake_sys_path[0]).resolve(), _REPO_ROOT.resolve()) + + + def test_non_transformer_head_with_get_optim_groups_still_uses_custom_groups(self): + module = self._load_train_vla_module() + agent = FakeIMFAgent() + + optimizer = module.build_training_optimizer(agent, lr=1e-4, weight_decay=0.123) + + self.assertEqual(agent.noise_pred_net.optim_group_calls, [0.123]) + group_names = self._group_names(agent, optimizer) + self.assertEqual(group_names[0], {'noise_pred_net.proj.weight'}) + self.assertEqual(group_names[1], { + 'noise_pred_net.proj.bias', + 'noise_pred_net.norm.weight', + 'noise_pred_net.norm.bias', + }) + self.assertEqual(group_names[2], {'backbone.weight', 'backbone.bias', 'adapter.weight'}) + + + def test_transformer_training_prefers_head_optim_groups_and_keeps_remaining_trainable_params(self): + module = self._load_train_vla_module() + agent = FakeTransformerAgent() + cfg = self._make_cfg() + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=lambda *args, **kwargs: FakeLoader()), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \ + mock.patch.object(module, 'AdamW', RecordingAdamW), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable): + module.main(cfg) + finally: + os.chdir(previous_cwd) + + self.assertEqual(agent.noise_pred_net.optim_group_calls, [cfg.train.weight_decay]) + + optimizer = RecordingAdamW.created[-1] + trainable_names = { + name for name, param in agent.named_parameters() if param.requires_grad + } + grouped_names = self._group_names(agent, optimizer) + optimizer_names = set().union(*grouped_names) + expected_head_names = { + 'noise_pred_net.proj.weight', + 'noise_pred_net.proj.bias', + 'noise_pred_net.norm.weight', + 'noise_pred_net.norm.bias', + } + expected_non_head_names = { + 'backbone.weight', + 'backbone.bias', + 'adapter.weight', + } + + self.assertEqual(grouped_names[0], {'noise_pred_net.proj.weight'}) + self.assertEqual(grouped_names[1], expected_head_names - {'noise_pred_net.proj.weight'}) + self.assertEqual(grouped_names[2], expected_non_head_names) + self.assertEqual(optimizer.param_groups[0]['weight_decay'], cfg.train.weight_decay) + self.assertEqual(optimizer.param_groups[1]['weight_decay'], 0.0) + self.assertEqual(optimizer.param_groups[2]['weight_decay'], cfg.train.weight_decay) + self.assertEqual(optimizer_names, trainable_names) + + flattened_param_ids = [ + id(param) + for group in optimizer.param_groups + for param in group['params'] + ] + self.assertEqual(len(flattened_param_ids), len(set(flattened_param_ids))) + self.assertNotIn('frozen.weight', optimizer_names) + self.assertNotIn('frozen.bias', optimizer_names) + + def test_any_head_with_get_optim_groups_uses_custom_groups_even_without_transformer_head_type(self): + module = self._load_train_vla_module() + agent = FakeTransformerAgent(head_type='imf') + + with mock.patch.object(module, 'AdamW', RecordingAdamW): + optimizer = module.build_training_optimizer(agent, lr=1e-4, weight_decay=0.123) + + self.assertEqual(agent.noise_pred_net.optim_group_calls, [0.123]) + grouped_names = self._group_names(agent, optimizer) + self.assertEqual(grouped_names[0], {'noise_pred_net.proj.weight'}) + self.assertEqual( + grouped_names[1], + {'noise_pred_net.proj.bias', 'noise_pred_net.norm.weight', 'noise_pred_net.norm.bias'}, + ) + self.assertEqual(grouped_names[2], {'backbone.weight', 'backbone.bias', 'adapter.weight'}) + + def test_transformer_optimizer_ignores_frozen_head_params_returned_by_head_groups(self): + module = self._load_train_vla_module() + agent = FakeTransformerAgent() + agent.noise_pred_net.norm.bias.requires_grad = False + cfg = self._make_cfg() + + def fake_instantiate(config_node, **_kwargs): + if config_node is cfg.data: + return FakeDataset() + if config_node is cfg.agent: + return agent + raise AssertionError(f'unexpected instantiate config: {config_node!r}') + + with tempfile.TemporaryDirectory() as tempdir: + previous_cwd = os.getcwd() + try: + os.chdir(tempdir) + with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \ + mock.patch.object(module, 'DataLoader', side_effect=lambda *args, **kwargs: FakeLoader()), \ + mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \ + mock.patch.object(module, 'AdamW', RecordingAdamW), \ + mock.patch.object(module.torch, 'save', return_value=None), \ + mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable): + module.main(cfg) + finally: + os.chdir(previous_cwd) + + optimizer = RecordingAdamW.created[-1] + optimizer_names = set().union(*self._group_names(agent, optimizer)) + trainable_names = { + name for name, param in agent.named_parameters() if param.requires_grad + } + + self.assertEqual(agent.noise_pred_net.optim_group_calls, [cfg.train.weight_decay]) + self.assertEqual(optimizer_names, trainable_names) + self.assertNotIn('noise_pred_net.norm.bias', optimizer_names) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_transformer1d_external_alignment.py b/tests/test_transformer1d_external_alignment.py new file mode 100644 index 0000000..f3b199c --- /dev/null +++ b/tests/test_transformer1d_external_alignment.py @@ -0,0 +1,262 @@ +import contextlib +import importlib.util +import inspect +import sys +import types +import unittest +import warnings +from pathlib import Path + +import torch + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_LOCAL_MODULE_PATH = _REPO_ROOT / 'roboimi/vla/models/heads/transformer1d.py' +_EXTERNAL_CHECKOUT_ROOT = _REPO_ROOT.parent / 'diffusion_policy' +_TRANSFORMER_WARNING_MESSAGE = ( + r'enable_nested_tensor is True, but self.use_nested_tensor is False ' + r'because encoder_layer\.norm_first was True' +) +_MISSING = object() + + +def _load_module_from_path(name: str, path: Path, *, register: bool = False): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + if register: + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _resolve_external_module_paths(external_checkout_root: Path): + diffusion_policy_root = external_checkout_root / 'diffusion_policy' + paths = { + 'positional_embedding': diffusion_policy_root / 'model/diffusion/positional_embedding.py', + 'module_attr_mixin': diffusion_policy_root / 'model/common/module_attr_mixin.py', + 'transformer_for_diffusion': diffusion_policy_root / 'model/diffusion/transformer_for_diffusion.py', + } + if not all(path.exists() for path in paths.values()): + return None + return paths + + +@contextlib.contextmanager +def _temporary_registered_modules(): + previous_modules = {} + + def remember(name: str) -> None: + if name not in previous_modules: + previous_modules[name] = sys.modules.get(name, _MISSING) + + def ensure_package(name: str) -> None: + if not name or name in sys.modules: + return + remember(name) + package = types.ModuleType(name) + package.__path__ = [] + sys.modules[name] = package + + def load(name: str, path: Path): + package_parts = name.split('.')[:-1] + for idx in range(1, len(package_parts) + 1): + ensure_package('.'.join(package_parts[:idx])) + + remember(name) + return _load_module_from_path(name, path, register=True) + + try: + yield load + finally: + for name, previous in reversed(list(previous_modules.items())): + if previous is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + + +@contextlib.contextmanager +def _suppress_nested_tensor_warning(): + with warnings.catch_warnings(): + warnings.filterwarnings( + 'ignore', + message=_TRANSFORMER_WARNING_MESSAGE, + category=UserWarning, + module=r'torch\.nn\.modules\.transformer', + ) + yield + + +def _load_local_module(): + return _load_module_from_path('local_transformer1d_alignment', _LOCAL_MODULE_PATH) + + +class Transformer1DExternalAlignmentTest(unittest.TestCase): + def _load_transformer_classes_or_skip(self): + external_paths = _resolve_external_module_paths(_EXTERNAL_CHECKOUT_ROOT) + if external_paths is None: + self.skipTest(f'external diffusion_policy checkout unavailable under {_EXTERNAL_CHECKOUT_ROOT}') + + local_module = _load_local_module() + with _temporary_registered_modules() as load_external: + load_external( + 'diffusion_policy.model.diffusion.positional_embedding', + external_paths['positional_embedding'], + ) + load_external( + 'diffusion_policy.model.common.module_attr_mixin', + external_paths['module_attr_mixin'], + ) + external_module = load_external( + 'diffusion_policy.model.diffusion.transformer_for_diffusion', + external_paths['transformer_for_diffusion'], + ) + + return local_module.Transformer1D, local_module.create_transformer1d, external_module.TransformerForDiffusion + + def _optim_group_names(self, model, groups): + names_by_param = {id(param): name for name, param in model.named_parameters()} + return [ + {names_by_param[id(param)] for param in group['params']} + for group in groups + ] + + def test_missing_external_checkout_resolution_returns_none(self): + self.assertIsNone(_resolve_external_module_paths(_REPO_ROOT / '__missing_diffusion_policy_checkout__')) + + def test_external_loader_restores_injected_sys_modules(self): + external_paths = _resolve_external_module_paths(_EXTERNAL_CHECKOUT_ROOT) + if external_paths is None: + self.skipTest(f'external diffusion_policy checkout unavailable under {_EXTERNAL_CHECKOUT_ROOT}') + + watched_names = [ + 'diffusion_policy', + 'diffusion_policy.model', + 'diffusion_policy.model.common', + 'diffusion_policy.model.common.module_attr_mixin', + 'diffusion_policy.model.diffusion', + 'diffusion_policy.model.diffusion.positional_embedding', + 'diffusion_policy.model.diffusion.transformer_for_diffusion', + ] + before = {name: sys.modules.get(name, _MISSING) for name in watched_names} + + with _temporary_registered_modules() as load_external: + load_external( + 'diffusion_policy.model.diffusion.positional_embedding', + external_paths['positional_embedding'], + ) + load_external( + 'diffusion_policy.model.common.module_attr_mixin', + external_paths['module_attr_mixin'], + ) + load_external( + 'diffusion_policy.model.diffusion.transformer_for_diffusion', + external_paths['transformer_for_diffusion'], + ) + + after = {name: sys.modules.get(name, _MISSING) for name in watched_names} + self.assertEqual(after, before) + + def test_transformer1d_preserves_local_direct_call_defaults(self): + local_module = _load_local_module() + ctor = inspect.signature(local_module.Transformer1D.__init__).parameters + helper = inspect.signature(local_module.create_transformer1d).parameters + + self.assertEqual(ctor['n_layer'].default, 8) + self.assertEqual(ctor['n_head'].default, 8) + self.assertEqual(ctor['n_emb'].default, 256) + self.assertEqual(helper['n_layer'].default, 8) + self.assertEqual(helper['n_head'].default, 8) + self.assertEqual(helper['n_emb'].default, 256) + + def test_time_as_cond_false_token_accounting_matches_external(self): + Transformer1D, _, TransformerForDiffusion = self._load_transformer_classes_or_skip() + self.assertIn('time_as_cond', inspect.signature(Transformer1D.__init__).parameters) + + config = dict( + input_dim=4, + output_dim=4, + horizon=6, + n_obs_steps=3, + cond_dim=0, + n_layer=2, + n_head=2, + n_emb=8, + p_drop_emb=0.0, + p_drop_attn=0.0, + causal_attn=False, + time_as_cond=False, + obs_as_cond=False, + n_cond_layers=0, + ) + + torch.manual_seed(5) + with _suppress_nested_tensor_warning(): + external_model = TransformerForDiffusion(**config) + local_model = Transformer1D(**config) + external_model.eval() + local_model.eval() + + self.assertEqual(local_model.T, external_model.T) + self.assertEqual(local_model.T_cond, external_model.T_cond) + self.assertEqual(local_model.time_as_cond, external_model.time_as_cond) + self.assertEqual(local_model.obs_as_cond, external_model.obs_as_cond) + self.assertEqual(local_model.encoder_only, external_model.encoder_only) + + def test_nocausal_state_dict_forward_and_optim_groups_match_external(self): + Transformer1D, _, TransformerForDiffusion = self._load_transformer_classes_or_skip() + config = dict( + input_dim=4, + output_dim=4, + horizon=6, + n_obs_steps=3, + cond_dim=5, + n_layer=2, + n_head=2, + n_emb=8, + p_drop_emb=0.0, + p_drop_attn=0.0, + causal_attn=False, + obs_as_cond=True, + n_cond_layers=1, + ) + + torch.manual_seed(7) + with _suppress_nested_tensor_warning(): + external_model = TransformerForDiffusion(**config) + local_model = Transformer1D(**config) + external_model.eval() + local_model.eval() + + external_state_dict = external_model.state_dict() + self.assertEqual(set(local_model.state_dict().keys()), set(external_state_dict.keys())) + local_model.load_state_dict(external_state_dict, strict=True) + + batch_size = 2 + sample = torch.randn(batch_size, config['horizon'], config['input_dim']) + cond = torch.randn(batch_size, config['n_obs_steps'], config['cond_dim']) + timestep = torch.tensor([11, 17], dtype=torch.long) + + with torch.no_grad(): + external_out = external_model(sample=sample, timestep=timestep, cond=cond) + local_out = local_model(sample=sample, timestep=timestep, cond=cond) + + self.assertEqual(local_out.shape, (batch_size, config['horizon'], config['output_dim'])) + self.assertEqual(local_out.shape, external_out.shape) + self.assertTrue(torch.allclose(local_out, external_out, atol=1e-6, rtol=1e-5)) + + weight_decay = 0.123 + external_groups = external_model.get_optim_groups(weight_decay=weight_decay) + local_groups = local_model.get_optim_groups(weight_decay=weight_decay) + + self.assertEqual(len(local_groups), len(external_groups)) + self.assertEqual([group['weight_decay'] for group in local_groups], [weight_decay, 0.0]) + self.assertEqual( + self._optim_group_names(local_model, local_groups), + self._optim_group_names(external_model, external_groups), + ) + + +if __name__ == '__main__': + unittest.main()