diff --git a/README.md b/README.md index cf3a1b381e16a8a3eac4c15c9c4f09e1058558c9..b7e7a115a153c44149906a7a519d88688cd2e452 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ This package contains the server-authoritative Florence–Bridge OPTR code, its ## What is included +- [`real_robot/`](real_robot/README.md)：Xtrainer2 真机数据转换、训练、双臂动作适配和 policy server 源码，以及中文交接说明。 - `common_start_20k/checkpoints/steps_20000_accelerate/`: **full resumable state** (model, optimizer, scheduler metadata and per-rank RNG state). Start every new continuation arm here. - `checkpoints/uniform/steps_20500_pytorch_model.pt`: model-only Uniform +500 snapshot. - `checkpoints/residual/steps_20500_pytorch_model.pt`: model-only Residual-OPTR +500 snapshot. diff --git a/real_robot/AGENTS.md b/real_robot/AGENTS.md new file mode 100644 index 0000000000000000000000000000000000000000..8cf37f2b7426ef460fde8db57e2a84853d2430a9 --- /dev/null +++ b/real_robot/AGENTS.md @@ -0,0 +1,10 @@ +# Xtrainer2 真机源码交接 + +本目录用于真机数据处理、训练与 policy server 部署。先阅读 `README.md` 和 `SOURCE_FILES.json`，从 `source/` 进入原始源码根目录。 + +- 使用本目录内的 `starVLA`、`retime` 和 `deployment` 包。 +- 按所选真机配置核对 GPU 数、global batch、训练预算、数据混合和恢复 step。 +- 模型权重、`config.yaml` 和 `dataset_statistics.json` 成套部署。 +- 相机顺序为 high、left wrist、right wrist，机器人侧 16 维 xyzw 布局，模型内部 20 维 rotation-6D。 +- 硬件驱动和执行循环接入实际机械臂客户端，执行动作前核对位置单位、夹爪约定、左右臂顺序和控制频率。 +- 源码及来源文件保留上游内容，新增实验使用独立输出目录。 diff --git a/real_robot/README.md b/real_robot/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c2bb46fec9b2081df096c1776621e931a3cb7215 --- /dev/null +++ b/real_robot/README.md @@ -0,0 +1,72 @@ +# Xtrainer2 真机代码交接 + +这份源码来自你的 GitHub 仓库，整理在 HF 的独立目录 `real_robot/`。训练与推理服务以 `ReTime_source/main` 的 Xtrainer2 实现为主体，原开发仓库 `OPTR-StarVLA/main` 提供 HDF5 数据转换与 LeRobot 格式验证工具。逐文件来源见 [SOURCE_FILES.json](SOURCE_FILES.json)。 + +## 代码入口 + +| 用途 | 文件 | +|---|---| +| HDF5 → LeRobot 数据转换 | `source/scripts/convert_xtrainer2_hdf5_to_lerobot.py` | +| 数据格式与四元数约定验证 | `source/scripts/validate_xtrainer2_lerobot.py` | +| 数据注册、三相机及双臂状态配置 | `source/examples/Xtrainer2/train_files/data_registry/data_config.py` | +| 数据 support 构建 | `source/examples/Xtrainer2/train_files/build_pack_phone_support.py` 和 `build_pick_bottle_support.py` | +| 训练启动 | `source/examples/Xtrainer2/train_files/run_xtrainer2_train.sh` | +| ReTime proposal 与 continuation | `source/retime/cli.py` 和 `source/scripts/run_xtrainer2_arm.sh` | +| 真机 policy server 启动 | `source/examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh` | +| 双臂动作与四元数转换 | `source/deployment/model_server/xtrainer2_command.py` | +| WebSocket 客户端通信工具 | `source/deployment/model_server/tools/websocket_policy_client.py` | + +`source/` 保留上游包布局。进入该目录再运行命令，使 `starVLA`、`retime` 和 `deployment` 使用这份源码。 + +## 环境与依赖 + +目标运行环境为 Linux、Python ≥ 3.10。原依赖列表为 [requirements.txt](source/requirements.txt) 和 [requirements-xtrainer2.txt](source/requirements-xtrainer2.txt)，其中 PyTorch 为 2.7.1，Transformers 为 4.57.0，Accelerate 为 1.13.0。 + +双臂四元数与 rotation-6D 转换需要匹配 Python、PyTorch、CUDA 的 PyTorch3D。数据转换另需 `h5py`、`pandas`、Pillow 和可执行的 `ffmpeg`。按目标 GPU 环境安装相应依赖。 + +## 数据与训练配置 + +训练配置放在 `source/configs/`，包含 pack-phone、pick-bottle、smoke 和从 pack-phone common start 延续的配置。原配置使用 Florence、8 GPU、每卡 micro-batch 4、累积 4，global batch 为 128。pack-phone 的完整训练预算为 52,380 updates，warmup 为 5,000。continuation 配置从 `steps_17460_accelerate` 延续，训练 seed 为 117465，基础训练 seed 为 42。 + +运行前按配置设置 `FLORENCE_MODEL_PATH`、`XTRAINER2_DATA_ROOT`、`OPTR_RUN_ROOT` 及数据 support 路径。训练入口会检查配置实际引用的环境变量。具体实现见 [Xtrainer2 数据说明](source/examples/Xtrainer2/README.md)。 + +## 推理服务 + +准备真机 checkpoint 对应的运行目录，模型、配置、归一化统计保持成套。 + +```text +your_real_robot_run/ +├── config.yaml +├── dataset_statistics.json +└── checkpoints/ + └── your_model.pt +``` + +在下载后的 `real_robot/source` 中执行以下命令，将示例路径改为实际路径。 + +```bash +export OPTR_PYTHON=/path/to/your/python +export FLORENCE_MODEL_PATH=/path/to/Florence-2-large +export TMPDIR="$PWD/runtime_tmp" +mkdir -p "$TMPDIR" + +bash examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh \ + /path/to/your_real_robot_run/checkpoints/your_model.pt 0 10093 +``` + +服务使用 BF16，并在连接握手中返回动作维数、相机顺序、四元数约定和 action chunk 长度。 + +## 机械臂客户端接口 + +相机顺序为 `cam_high`、`cam_left_wrist`、`cam_right_wrist`，图像使用 HWC RGB uint8。客户端提供 `image`、`lang` 和原始 16 维 `state`。双臂状态与动作按每臂位置 3 维、四元数 4 维、夹爪 1 维排列，机器人侧四元数顺序为 `xyzw`。服务端负责归一化和坐标表示转换，模型内部使用 20 维 rotation-6D 表示，返回动作使用机器人侧的 16 维布局。 + +pick-bottle 的相对末端动作由服务端结合当前状态还原为绝对动作。真实机械臂驱动、相机采集和执行循环由独立硬件客户端完成，需接入师兄实际使用的控制程序。上游接口说明见 [ROBOT.md](source/examples/Xtrainer2/ROBOT.md)。 + +## 本次交接验证 + +- 129 个源文件与来源清单的字节数一致，逐文件来源见 `SOURCE_FILES.json`。 +- 95 个 Python 文件通过语法检查，7 个 YAML、1 个 JSON 和 1 个 TOML 文件通过解析检查。 +- 8 个 shell 入口通过 `bash -n`，`python -m retime --help` 可正常运行。 +- 使用 jlu66 现有 Python 环境在 CPU 上执行 13 项配置测试，全部通过，覆盖双臂维数、global batch、梯度累积、ZeRO-3 权重保存设置、训练种子和 continuation 参数。 + +实际模型加载、GPU 推理和机械臂动作执行需在配套设备与 checkpoint 上验证。 diff --git a/real_robot/SOURCE_FILES.json b/real_robot/SOURCE_FILES.json new file mode 100644 index 0000000000000000000000000000000000000000..b83ea6f5a7ce3dccc409459c148bb3268bc58934 --- /dev/null +++ b/real_robot/SOURCE_FILES.json @@ -0,0 +1,909 @@ +{ + "schema_version": 1, + "description": "Xtrainer2 real-robot training and policy-serving source handoff", + "files": [ + { + "path": "source/configs/deepspeed_florence_8gpu.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/deepspeed_florence_8gpu.yaml", + "bytes": 702 + }, + { + "path": "source/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml", + "bytes": 2837 + }, + { + "path": "source/configs/florence_xtrainer2_pack_phone_residual_resume_smoke.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/florence_xtrainer2_pack_phone_residual_resume_smoke.yaml", + "bytes": 2967 + }, + { + "path": "source/configs/florence_xtrainer2_pack_phone_smoke.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/florence_xtrainer2_pack_phone_smoke.yaml", + "bytes": 3954 + }, + { + "path": "source/configs/florence_xtrainer2_pack_phone_uniform.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/florence_xtrainer2_pack_phone_uniform.yaml", + "bytes": 3898 + }, + { + "path": "source/configs/florence_xtrainer2_pick_bottle_rel_aug_smoke.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/florence_xtrainer2_pick_bottle_rel_aug_smoke.yaml", + "bytes": 3703 + }, + { + "path": "source/configs/florence_xtrainer2_pick_bottle_rel_aug_uniform.yaml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "configs/florence_xtrainer2_pick_bottle_rel_aug_uniform.yaml", + "bytes": 4104 + }, + { + "path": "source/deployment/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/__init__.py", + "bytes": 0 + }, + { + "path": "source/deployment/model_server/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/__init__.py", + "bytes": 0 + }, + { + "path": "source/deployment/model_server/policy_norm_processor.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/policy_norm_processor.py", + "bytes": 22282 + }, + { + "path": "source/deployment/model_server/policy_wrapper.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/policy_wrapper.py", + "bytes": 12974 + }, + { + "path": "source/deployment/model_server/server_policy.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/server_policy.py", + "bytes": 2741 + }, + { + "path": "source/deployment/model_server/tools/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/tools/__init__.py", + "bytes": 0 + }, + { + "path": "source/deployment/model_server/tools/image_tools.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/tools/image_tools.py", + "bytes": 4765 + }, + { + "path": "source/deployment/model_server/tools/msgpack_numpy.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/tools/msgpack_numpy.py", + "bytes": 2016 + }, + { + "path": "source/deployment/model_server/tools/websocket_policy_client.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/tools/websocket_policy_client.py", + "bytes": 2728 + }, + { + "path": "source/deployment/model_server/tools/websocket_policy_server.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/tools/websocket_policy_server.py", + "bytes": 5813 + }, + { + "path": "source/deployment/model_server/xtrainer2_command.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "deployment/model_server/xtrainer2_command.py", + "bytes": 3595 + }, + { + "path": "source/examples/Xtrainer2/INFERENCE.md", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/INFERENCE.md", + "bytes": 180 + }, + { + "path": "source/examples/Xtrainer2/README.md", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/README.md", + "bytes": 1511 + }, + { + "path": "source/examples/Xtrainer2/ROBOT.md", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/ROBOT.md", + "bytes": 1124 + }, + { + "path": "source/examples/Xtrainer2/train_files/build_pack_phone_support.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/train_files/build_pack_phone_support.py", + "bytes": 3018 + }, + { + "path": "source/examples/Xtrainer2/train_files/build_pick_bottle_support.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/train_files/build_pick_bottle_support.py", + "bytes": 3071 + }, + { + "path": "source/examples/Xtrainer2/train_files/data_registry/data_config.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/train_files/data_registry/data_config.py", + "bytes": 6348 + }, + { + "path": "source/examples/Xtrainer2/train_files/modality.json", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/train_files/modality.json", + "bytes": 1877 + }, + { + "path": "source/examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh", + "bytes": 4398 + }, + { + "path": "source/examples/Xtrainer2/train_files/run_xtrainer2_train.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "examples/Xtrainer2/train_files/run_xtrainer2_train.sh", + "bytes": 3016 + }, + { + "path": "source/LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "LICENSE", + "bytes": 1387 + }, + { + "path": "source/licenses/Diffusion-Policy-LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "licenses/Diffusion-Policy-LICENSE", + "bytes": 1106 + }, + { + "path": "source/licenses/Isaac-GR00T-LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "licenses/Isaac-GR00T-LICENSE", + "bytes": 10808 + }, + { + "path": "source/licenses/MODIFIED_FILES.txt", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "licenses/MODIFIED_FILES.txt", + "bytes": 765 + }, + { + "path": "source/licenses/msgpack-numpy-LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "licenses/msgpack-numpy-LICENSE", + "bytes": 1551 + }, + { + "path": "source/licenses/OpenVLA-LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "licenses/OpenVLA-LICENSE", + "bytes": 1105 + }, + { + "path": "source/licenses/X-VLA-LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "licenses/X-VLA-LICENSE", + "bytes": 11357 + }, + { + "path": "source/NOTICE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "NOTICE", + "bytes": 2060 + }, + { + "path": "source/pyproject.toml", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "pyproject.toml", + "bytes": 1790 + }, + { + "path": "source/README.md", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "README.md", + "bytes": 10922 + }, + { + "path": "source/requirements-qwen.txt", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "requirements-qwen.txt", + "bytes": 269 + }, + { + "path": "source/requirements-xtrainer2.txt", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "requirements-xtrainer2.txt", + "bytes": 440 + }, + { + "path": "source/requirements.txt", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "requirements.txt", + "bytes": 574 + }, + { + "path": "source/retime/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/__init__.py", + "bytes": 160 + }, + { + "path": "source/retime/__main__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/__main__.py", + "bytes": 79 + }, + { + "path": "source/retime/checkpoint.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/checkpoint.py", + "bytes": 864 + }, + { + "path": "source/retime/cli.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/cli.py", + "bytes": 13905 + }, + { + "path": "source/retime/commands/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/commands/__init__.py", + "bytes": 43 + }, + { + "path": "source/retime/commands/proposal.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/commands/proposal.py", + "bytes": 25822 + }, + { + "path": "source/retime/commands/shuffle.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/commands/shuffle.py", + "bytes": 1482 + }, + { + "path": "source/retime/objective.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/objective.py", + "bytes": 6506 + }, + { + "path": "source/retime/prepare.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/prepare.py", + "bytes": 1145 + }, + { + "path": "source/retime/pretrain.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/pretrain.py", + "bytes": 7183 + }, + { + "path": "source/retime/proposal.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/proposal.py", + "bytes": 9294 + }, + { + "path": "source/retime/sampler.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/sampler.py", + "bytes": 13229 + }, + { + "path": "source/retime/scoring/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/__init__.py", + "bytes": 218 + }, + { + "path": "source/retime/scoring/cheap_proposals.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/cheap_proposals.py", + "bytes": 7068 + }, + { + "path": "source/retime/scoring/cnf_likelihood.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/cnf_likelihood.py", + "bytes": 10056 + }, + { + "path": "source/retime/scoring/discrete_flow_likelihood.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/discrete_flow_likelihood.py", + "bytes": 7055 + }, + { + "path": "source/retime/scoring/flow_residual.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/flow_residual.py", + "bytes": 9357 + }, + { + "path": "source/retime/scoring/partial_gradient.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/partial_gradient.py", + "bytes": 4067 + }, + { + "path": "source/retime/scoring/residual_cache.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/residual_cache.py", + "bytes": 2815 + }, + { + "path": "source/retime/scoring/shuffle.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/shuffle.py", + "bytes": 4229 + }, + { + "path": "source/retime/scoring/starvla_likelihood.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/starvla_likelihood.py", + "bytes": 5643 + }, + { + "path": "source/retime/scoring/utility_audit.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "retime/scoring/utility_audit.py", + "bytes": 18264 + }, + { + "path": "source/scripts/build_residual_temporal_proposal.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/build_residual_temporal_proposal.py", + "bytes": 243 + }, + { + "path": "source/scripts/build_shuffle_temporal_proposal.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/build_shuffle_temporal_proposal.py", + "bytes": 252 + }, + { + "path": "source/scripts/convert_xtrainer2_hdf5_to_lerobot.py", + "repository": "miaochengy-pngu/OPTR-StarVLA", + "ref": "main", + "source_path": "scripts/convert_xtrainer2_hdf5_to_lerobot.py", + "bytes": 26880 + }, + { + "path": "source/scripts/launch_pack_phone_optr_proposal.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/launch_pack_phone_optr_proposal.sh", + "bytes": 247 + }, + { + "path": "source/scripts/launch_pack_phone_residual_20ep_to_60ep.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/launch_pack_phone_residual_20ep_to_60ep.sh", + "bytes": 2187 + }, + { + "path": "source/scripts/launch_pack_phone_uniform.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/launch_pack_phone_uniform.sh", + "bytes": 1885 + }, + { + "path": "source/scripts/launch_pick_bottle_rel_aug_smoke.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/launch_pick_bottle_rel_aug_smoke.sh", + "bytes": 1156 + }, + { + "path": "source/scripts/launch_pick_bottle_rel_aug_uniform.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/launch_pick_bottle_rel_aug_uniform.sh", + "bytes": 1441 + }, + { + "path": "source/scripts/run_xtrainer2_arm.sh", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "scripts/run_xtrainer2_arm.sh", + "bytes": 6276 + }, + { + "path": "source/scripts/validate_xtrainer2_lerobot.py", + "repository": "miaochengy-pngu/OPTR-StarVLA", + "ref": "main", + "source_path": "scripts/validate_xtrainer2_lerobot.py", + "bytes": 5871 + }, + { + "path": "source/starVLA/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/__init__.py", + "bytes": 0 + }, + { + "path": "source/starVLA/dataloader/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/__init__.py", + "bytes": 4672 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/__init__.py", + "bytes": 0 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/data_config.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/data_config.py", + "bytes": 35501 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/datasets.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/datasets.py", + "bytes": 144422 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/embodiment_tags.py", + "bytes": 2459 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/mixtures.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/mixtures.py", + "bytes": 19387 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/registry.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/registry.py", + "bytes": 5960 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/schema.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/schema.py", + "bytes": 8590 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/transform/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/transform/__init__.py", + "bytes": 1261 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/transform/base.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/transform/base.py", + "bytes": 5154 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/transform/concat.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/transform/concat.py", + "bytes": 9371 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/transform/state_action.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/transform/state_action.py", + "bytes": 27803 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/transform/video.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/transform/video.py", + "bytes": 24512 + }, + { + "path": "source/starVLA/dataloader/gr00t_lerobot/video.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/gr00t_lerobot/video.py", + "bytes": 14043 + }, + { + "path": "source/starVLA/dataloader/infoflow_sampler.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/infoflow_sampler.py", + "bytes": 180 + }, + { + "path": "source/starVLA/dataloader/lerobot_datasets.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/lerobot_datasets.py", + "bytes": 6050 + }, + { + "path": "source/starVLA/dataloader/relative_pose.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/relative_pose.py", + "bytes": 3359 + }, + { + "path": "source/starVLA/dataloader/temporal_proposal.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/dataloader/temporal_proposal.py", + "bytes": 181 + }, + { + "path": "source/starVLA/model/framework/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/framework/__init__.py", + "bytes": 202 + }, + { + "path": "source/starVLA/model/framework/base_framework.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/framework/base_framework.py", + "bytes": 11351 + }, + { + "path": "source/starVLA/model/framework/share_tools.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/framework/share_tools.py", + "bytes": 21898 + }, + { + "path": "source/starVLA/model/framework/VLM4A/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/framework/VLM4A/__init__.py", + "bytes": 296 + }, + { + "path": "source/starVLA/model/framework/VLM4A/QwenGR00T.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/framework/VLM4A/QwenGR00T.py", + "bytes": 15342 + }, + { + "path": "source/starVLA/model/modules/action_model/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/action_model/__init__.py", + "bytes": 122 + }, + { + "path": "source/starVLA/model/modules/action_model/flow_matching_head/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/action_model/flow_matching_head/__init__.py", + "bytes": 679 + }, + { + "path": "source/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/action_model/flow_matching_head/action_encoder.py", + "bytes": 3431 + }, + { + "path": "source/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py", + "bytes": 14773 + }, + { + "path": "source/starVLA/model/modules/action_model/GR00T_ActionHeader.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/action_model/GR00T_ActionHeader.py", + "bytes": 21024 + }, + { + "path": "source/starVLA/model/modules/vlm/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/vlm/__init__.py", + "bytes": 481 + }, + { + "path": "source/starVLA/model/modules/vlm/Florence2.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/vlm/Florence2.py", + "bytes": 12888 + }, + { + "path": "source/starVLA/model/modules/vlm/QWen3.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/modules/vlm/QWen3.py", + "bytes": 7253 + }, + { + "path": "source/starVLA/model/tools.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/model/tools.py", + "bytes": 17659 + }, + { + "path": "source/starVLA/scoring/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/__init__.py", + "bytes": 209 + }, + { + "path": "source/starVLA/scoring/cheap_proposals.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/cheap_proposals.py", + "bytes": 196 + }, + { + "path": "source/starVLA/scoring/cnf_likelihood.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/cnf_likelihood.py", + "bytes": 195 + }, + { + "path": "source/starVLA/scoring/discrete_flow_likelihood.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/discrete_flow_likelihood.py", + "bytes": 205 + }, + { + "path": "source/starVLA/scoring/flow_residual.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/flow_residual.py", + "bytes": 194 + }, + { + "path": "source/starVLA/scoring/partial_gradient.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/partial_gradient.py", + "bytes": 197 + }, + { + "path": "source/starVLA/scoring/residual_cache.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/residual_cache.py", + "bytes": 195 + }, + { + "path": "source/starVLA/scoring/shuffle.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/shuffle.py", + "bytes": 188 + }, + { + "path": "source/starVLA/scoring/starvla_likelihood.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/starvla_likelihood.py", + "bytes": 199 + }, + { + "path": "source/starVLA/scoring/utility_audit.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/scoring/utility_audit.py", + "bytes": 194 + }, + { + "path": "source/starVLA/training/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/__init__.py", + "bytes": 0 + }, + { + "path": "source/starVLA/training/infoflow.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/infoflow.py", + "bytes": 182 + }, + { + "path": "source/starVLA/training/train_starvla.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/train_starvla.py", + "bytes": 99854 + }, + { + "path": "source/starVLA/training/trainer_utils/__init__.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/trainer_utils/__init__.py", + "bytes": 44 + }, + { + "path": "source/starVLA/training/trainer_utils/config_tracker.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/trainer_utils/config_tracker.py", + "bytes": 21344 + }, + { + "path": "source/starVLA/training/trainer_utils/monkey_patch.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/trainer_utils/monkey_patch.py", + "bytes": 19353 + }, + { + "path": "source/starVLA/training/trainer_utils/overwatch.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/trainer_utils/overwatch.py", + "bytes": 6848 + }, + { + "path": "source/starVLA/training/trainer_utils/trainer_tools.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "starVLA/training/trainer_utils/trainer_tools.py", + "bytes": 23838 + }, + { + "path": "source/tests/test_xtrainer2_optr_config.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "tests/test_xtrainer2_optr_config.py", + "bytes": 6916 + }, + { + "path": "source/tests/test_xtrainer2_pack_phone_config.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "tests/test_xtrainer2_pack_phone_config.py", + "bytes": 11364 + }, + { + "path": "source/tests/test_xtrainer2_pick_bottle_config.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "tests/test_xtrainer2_pick_bottle_config.py", + "bytes": 4381 + }, + { + "path": "source/tests/test_xtrainer2_policy_server_contract.py", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "tests/test_xtrainer2_policy_server_contract.py", + "bytes": 9812 + }, + { + "path": "source/third_party/frameskip/LICENSE", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "third_party/frameskip/LICENSE", + "bytes": 1387 + }, + { + "path": "source/UPSTREAM.md", + "repository": "miaochengy-pngu/ReTime_source", + "ref": "main", + "source_path": "UPSTREAM.md", + "bytes": 1460 + } + ] +} diff --git a/real_robot/source/LICENSE b/real_robot/source/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..c636f73528f042c33c1e9a502a0f685316e44a11 --- /dev/null +++ b/real_robot/source/LICENSE @@ -0,0 +1,24 @@ + MIT License + + Copyright (c) StarVLA Team. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + Rebases are allowed for forks and feature branches. When rebasing from upstream StarVLA, use descriptive commit messages, e.g., "chore: clone from StarVLA". + Preserve attribution: keep at least the two latest upstream StarVLA commits as separate. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE diff --git a/real_robot/source/NOTICE b/real_robot/source/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..d58afa4485ab5c8364959f69233e5f7ca009cd13 --- /dev/null +++ b/real_robot/source/NOTICE @@ -0,0 +1,40 @@ +ReTime is a research extension of StarVLA, not an official StarVLA release. + +StarVLA: https://github.com/starVLA/starVLA +Base commit: 3422b9f2387b6f682cf02802904a77b23ab13afd +Copyright (c) StarVLA Team. +The complete upstream LICENSE is retained, including its additional attribution +and commit-preservation language. The original base commit and its preceding +upstream commit remain separate in the supplied Git history. + +Modifications in this artifact include temporal sampling/proposals, ordinary +importance correction, scoring utilities, model/data/trainer integration, +portable configurations, and the ReTime command-line interface. +Original ReTime additions in this artifact are distributed under the root +LICENSE terms. Their authors' identities are withheld during anonymous review. +Existing upstream author/copyright notices identify the respective third-party +contributors. + +NVIDIA Isaac-GR00T-derived code retains its source notices and Apache-2.0 terms. +See licenses/Isaac-GR00T-LICENSE. Modified source files carry a modification notice. +https://github.com/NVIDIA/Isaac-GR00T + +OpenVLA/Prismatic-derived logging code retains its attribution. +See licenses/OpenVLA-LICENSE (MIT). +https://github.com/openvla/openvla + +The Florence wrapper attributes X-VLA-derived code in its source. +See licenses/X-VLA-LICENSE (Apache-2.0). +https://github.com/2toinf/X-VLA + +The minimal FrameSkip source snapshot retains third_party/frameskip/LICENSE. +Rotation conversion code in starVLA/dataloader/gr00t_lerobot/transform/state_action.py +derives from Diffusion Policy. Its MIT license is retained in +licenses/Diffusion-Policy-LICENSE. +The adapted msgpack-numpy utility in deployment/model_server/tools/msgpack_numpy.py +retains the original BSD license in licenses/msgpack-numpy-LICENSE. +Third-party components remain subject to their respective licenses; the root +license does not replace these component-specific terms. + +No datasets, pretrained model weights, checkpoints, benchmark repositories, +or robot-client contents are redistributed in the working source tree. diff --git a/real_robot/source/README.md b/real_robot/source/README.md new file mode 100644 index 0000000000000000000000000000000000000000..67d9e6637c56d5dc05e675fb5922ccdab08e503a --- /dev/null +++ b/real_robot/source/README.md @@ -0,0 +1,222 @@ +# ReTime for StarVLA + +ReTime is an opt-in temporal-sampling extension for the accompanying StarVLA +backend. It preserves the dataset/trajectory distribution and changes the +within-trajectory temporal proposal: + +```text +q(t | e) = (1 - alpha) p(t | e) + alpha r(t | e) +training loss = mean((p / q) * per_sample_loss) +``` + +The weights are ordinary importance weights, not batch-normalized weights. +`K` controls the frozen-policy scoring grid. `shuffle` permutes the residual +masses within each trajectory; `uniform` uses the baseline draw stream. + +## Install + +Use Linux, Python 3.11, a CUDA driver compatible with PyTorch 2.7.1, and the +dependencies in `requirements.txt`. The core training packages are pinned to +the versions used in the GPU integration test. From this source tree: + +```bash +python -m pip install -r requirements.txt +python -m pip install -e . +retime --help +``` + +The supplied Qwen-GR00T configurations additionally use `AdamW8bit` and +`flash_attention_2`. Install their optional dependencies after the base +environment and select a Flash Attention build compatible with your CUDA and +PyTorch installation: + +```bash +python -m pip install -r requirements-qwen.txt +python -m pip install flash-attn==2.7.4.post1 --no-build-isolation +``` + +Flash Attention may require a local CUDA toolkit and build tools. The version +above was present in the GPU validation environments. Both production Qwen +continuation configurations use Flash Attention 2. The bounded Qwen Bridge +integration run used PyTorch SDPA on the same model/data path. + +This release supports the supplied Florence/Qwen GR00T configurations and +LeRobot data layouts. The backend includes the per-example loss and sampler +integration required by ReTime; it is not a claim of compatibility with every +unmodified StarVLA release or arbitrary policy architecture. + +## Training + +The trained common-start checkpoints used in the paper are not redistributed +with this anonymous artifact. Reviewers should either train their own +common-start policy using `retime-pretrain` or provide a compatible complete +checkpoint. The backbone, dataset, and benchmark assets must also be provided +locally. + +Select a configuration matching the policy architecture, action dimensions, +dataset mixture, and training budget. Matching supplied pairs are: + +| Setting | Policy configuration | Start → stop | Global batch | Accelerate configuration | +| --- | --- | ---: | ---: | --- | +| Florence × Bridge/RT-1 | `configs/florence_bridge_20k_plus500.yaml` | 20K → 20.5K | 64 | `configs/deepspeed_florence_2gpu.yaml` | +| Qwen × Bridge/RT-1 | `configs/qwen_bridge_29k_plus500.yaml` | 29K → 29.5K | 48 | `configs/qwen_single_1gpu.yaml` | +| Florence × RoboTwin | `configs/florence_robotwin_10k_plus500.yaml` | 10K → 10.5K | 64 | `configs/deepspeed_florence_robotwin_2gpu.yaml` | +| Qwen × RoboTwin | `configs/qwen_robotwin_11k_plus500.yaml` | 11K → 11.5K | 40 | `configs/deepspeed_qwen_robotwin_2gpu.yaml` | + +These are the four main simulation continuation contracts. Every arm restores +the complete shared starting state and performs exactly 500 optimizer updates. +Bridge runs additionally require `--statistics-path`. + +### Train the base policy locally + +Reviewers do not need the authors' original common-start checkpoint. +`retime-pretrain` can train a reviewer-owned common-start policy with the same +architecture, optimizer, scheduler, batch geometry, and dataset-sampling +contract as the selected continuation configuration. For Florence Bridge: + +```bash +retime-pretrain \ + --config configs/florence_bridge_20k_plus500.yaml \ + --accelerate-config configs/deepspeed_florence_2gpu.yaml \ + --base-steps 20000 \ + --vlm-path /path/to/Florence-2-large \ + --data-root /path/to/OXE_LEROBOT_DATASET \ + --statistics-path /path/to/dataset_statistics.json \ + --omit-data-identity-files \ + --output /path/to/florence-bridge-base +``` + +Use `--base-steps 29000`, `10000`, or `11000` with the other three configuration +pairs. The command builds and saves `manifest.json` before training. Two-process +configurations produce `checkpoints/steps__accelerate`; Qwen Bridge produces +`checkpoints/steps_29000_pytorch_model.pt` with adjacent training-state and +commit files. Both formats contain the model, optimizer, scheduler, RNG state, +sampling position, and completed-update counter required by exact continuation. + +`--omit-data-identity-files` stores the computed dataset inventory inside the +manifest. Use the same option for continuation. This avoids requiring separate +provenance and fingerprint files while still checking the dataset inventory +against the frozen manifest. + +For the Florence Bridge configuration: + +```bash +retime \ + --config configs/florence_bridge_20k_plus500.yaml \ + --accelerate-config configs/deepspeed_florence_2gpu.yaml \ + --base-policy /path/to/steps_20000_accelerate \ + --start resume \ + --manifest /path/to/florence-bridge-base/manifest.json \ + --vlm-path /path/to/Florence-2-large \ + --data-root /path/to/OXE_LEROBOT_DATASET \ + --statistics-path /path/to/dataset_statistics.json \ + --omit-data-identity-files \ + --alpha 0.2 --K 8 \ + --mode retime uniform shuffle \ + --output /path/to/new-run +``` + +Choose one mode or several. The command prepares valid-anchor support from the +dataset, generates the residual proposal once, constructs Shuffle when selected, +and trains each selected arm sequentially from the same base policy. A +uniform-only run does not score a proposal. Invalid probabilities, +incompatible tensor shapes, incomplete +checkpoint files, and mismatched data support still produce errors. + +`--base-policy` denotes the reviewer-provided common-start VLA checkpoint. +It is not bundled with this anonymous artifact and is not downloaded +automatically. `--vlm-path` denotes its Florence/Qwen backbone. They are +distinct inputs. The main configurations use +`--start resume`, which is required for the reported +500 continuation. +Florence Bridge and both RoboTwin settings accept a complete Accelerate +checkpoint directory. Qwen Bridge accepts its committed single-process model +file together with the adjacent training-state and commit files. + +Exact resume restores model, optimizer, scheduler, RNG, sampler position/data +order, and completed-update counter. Supply the `manifest.json` produced by +the base-policy run and keep its `config.full.yaml` in the original run layout. +When separate provenance/fingerprint files are available, omit +`--omit-data-identity-files` and set the corresponding environment variables in +both commands. A single-process model file is accepted only when its adjacent +`*_training_state.pt` and `*_checkpoint.json` files form a complete committed +checkpoint. + +`--start weights` remains available for a deliberate fresh optimizer/scheduler +run from model weights. It starts at update zero, so it does not reproduce any +of the four +500 experiments even when a continuation configuration is selected. + +The launcher never changes microbatch, world size, accumulation, optimizer, +schedule, or random seeds to select an arm. The supplied Bridge pairing remains +`4 x 2 x 8 = 64`. The table above supplies matching Accelerate configurations +for the other three settings. Select a smaller dataset mixture with `--data-mix` +when checking a subset; the default RoboTwin configurations cover all 50 tasks. + +Add `--dry-run` to validate inputs and print the command plan without writing or +launching. Successful runs produce `config.yaml`, `plan.json`, a proposal/cache, +and separate `runs/retime`, `runs/uniform`, and `runs/shuffle` directories. Use a +fresh output directory for each experiment. Existing proposal caches are bound +to the checkpoint content and support; incompatible caches cannot be reused. + +## Layout + +```text +retime/ method implementation and command-line entrypoint + pretrain.py reviewer-owned common-start policy training + objective.py ordinary p/q correction + sampler.py deterministic hierarchical draw stream + proposal.py full-support temporal proposal + scoring/ frozen-policy scoring and Shuffle control + commands/ proposal generation +starVLA/ attributed compatible training/model/data backend +scripts/ legacy entrypoints and optional analysis utilities +configs/ model, data, and training settings +licenses/ third-party license texts +``` + +Legacy `starVLA.*` method imports forward to `retime.*` for compatibility. +Analysis utilities are optional research tools and are not training prerequisites. + +## Compatibility + +| Interface | Support | +| --- | --- | +| StarVLA training | The original `starVLA/training/train_starvla.py` entrypoint is retained. | +| Legacy Python imports | Objective, sampler, proposal, and scoring modules forward to ReTime implementations. | +| Proposal and manifest files | Existing schemas and probability semantics are preserved. | +| Model initialization | Full model weights initialize a fresh optimizer and scheduler. | +| Full-state continuation | `--start resume` restores a complete Accelerate directory or committed single-process checkpoint. | +| Distributed training | Matching one- and two-GPU configurations are supplied for all four settings. | + +Florence-GR00T and Qwen-GR00T were exercised on GPUs with both Bridge/RT-1 and +RoboTwin data. Other policy architectures require a compatible per-sample loss +and dataset adapter. +Historical `OPTR_*` environment variables and `InfoFlow` module names in +example configurations are compatibility identifiers for existing datasets and +checkpoints. The supported command-line entrypoint for this artifact is `retime`. + +## Tests and scope + +```bash +python -m pip install pytest +python -m pytest tests/test_retime_plugin.py -q +``` + +These CPU tests check proposal generation, objective correction, deterministic +sampling, Shuffle invariants, legacy imports, CLI dispatch, and error propagation. +A separate GPU integration test completed proposal generation, two optimizer +updates, checkpoint saving, and strict model reload for each of the three modes +in all four settings. See `VALIDATION.md` for details. +Datasets, trained weights, checkpoints, and benchmark repositories are not +redistributed in the working source tree. + +## Attribution and anonymous review + +Built on [StarVLA](https://github.com/starVLA/starVLA), pinned to +`3422b9f2387b6f682cf02802904a77b23ab13afd`. See `UPSTREAM.md`, `LICENSE`, and +`NOTICE`. Retain the complete license, source copyright notices, and the +separate original upstream commits in the accompanying Git history. + +Third-party attribution is intentionally visible. The anonymous additions have +no author-account remote or personal development history. Publish through an +anonymous review endpoint, not an author's personal repository URL. Avoid +uploading local training outputs: they can contain machine paths and metadata. diff --git a/real_robot/source/UPSTREAM.md b/real_robot/source/UPSTREAM.md new file mode 100644 index 0000000000000000000000000000000000000000..9eae9cca3444e0aebb4e0d74073b6c5b447f0910 --- /dev/null +++ b/real_robot/source/UPSTREAM.md @@ -0,0 +1,29 @@ +# Upstream provenance + +- Repository: https://github.com/starVLA/starVLA +- Original branch: `starVLA` +- Base: `3422b9f2387b6f682cf02802904a77b23ab13afd` +- Preceding upstream commit: `61ea14279a22739316f9c1ca74c18d535247fe05` +- Base LICENSE SHA256: `6ca01970c4e66d2d5bebd0a00b55de5b4081017e0f0c97ddb415cda28e3e8ffa` + +The base revision was recovered from the original development snapshot's +provenance record and checked against the public upstream Git object and license. +The two original upstream commits are retained unchanged. A separate anonymous +commit contains the review artifact and modifications. Personal development +commits and remotes are not part of this distribution. + +The source tree is a selected, modified research backend, not an unmodified +checkout of that revision. Per-file notices and NOTICE describe the derivative +work. Do not replace LICENSE with a generic MIT template: this upstream text +contains additional attribution/commit-preservation language. + +The distribution ZIP includes `review-history.bundle`, containing the public +upstream ancestry and the anonymous review commit. Restore a Git repository with: + +```bash +git clone --branch anonymous-review review-history.bundle retime-review +``` + +This preserves the original upstream commit objects and creates no relationship +to the authors' personal development repository. Use the restored history when +publishing to an anonymous repository; do not squash away the upstream commits. diff --git a/real_robot/source/configs/deepspeed_florence_8gpu.yaml b/real_robot/source/configs/deepspeed_florence_8gpu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..92e2f8f429c37fd8fd073baf6efb4cdf14388c2c --- /dev/null +++ b/real_robot/source/configs/deepspeed_florence_8gpu.yaml @@ -0,0 +1,23 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + gradient_accumulation_steps: 4 + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: true + # Required: _save_accelerate_checkpoint() calls accelerator.get_state_dict(), + # which raises under ZeRO-3 unless the 16-bit weights are gathered on save. + # Without this the run trains fine and then dies at the first save_steps entry. + zero3_save_16bit_model: true + zero_stage: 3 +distributed_type: DEEPSPEED +downcast_bf16: "no" +enable_cpu_affinity: false +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +use_cpu: false diff --git a/real_robot/source/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml b/real_robot/source/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8008b260daad4b78ff823491369daa8bdde568d9 --- /dev/null +++ b/real_robot/source/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml @@ -0,0 +1,101 @@ +run_id: florence-xtrainer2-pack-phone-residual-20ep-to-60ep +run_root_dir: ${oc.env:OPTR_RUN_ROOT} +seed: 117465 +wandb_entity: null +wandb_project: OPTR-Florence-Xtrainer2 +wandb_group: pack-phone-optr +wandb_job_type: train +wandb_tags: [florence, xtrainer2, pack-phone, optr, residual, 20ep-common-start] +is_debug: false +version_id: "0.21" + +framework: + name: QwenGR00T + qwenvl: + base_vlm: ${oc.env:FLORENCE_MODEL_PATH} + attn_implementation: eager + vl_hidden_dim: 1024 + selected_view_indices: [0, 1, 2] + compile_vision_tower: false + action_model: + action_model_type: DiT-B + hidden_size: 1024 + add_pos_embed: true + max_seq_len: 1024 + action_dim: 20 + state_dim: 20 + action_horizon: 16 + future_action_window_size: 15 + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + num_inference_timesteps: 4 + num_target_vision_tokens: 32 + diffusion_model_cfg: + cross_attention_dim: 1024 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + norm_type: ada_norm + num_layers: 16 + output_dim: 1024 + positional_embeddings: null + +datasets: + vla_data: + dataset_py: lerobot_datasets + data_root_dir: ${oc.env:XTRAINER2_DATA_ROOT} + data_mix: xtrainer2_pack_phone + action_type: abs_eef + action_mode: abs + include_state: true + per_device_batch_size: 4 + image_size: [224, 224] + video_backend: pyav + num_workers: 8 + expected_world_size: 8 + traceable_sampling: true + dataset_provenance_path: ${oc.env:XTRAINER2_PACK_PHONE_PROVENANCE_PATH} + dataset_content_fingerprint_path: ${oc.env:XTRAINER2_PACK_PHONE_FINGERPRINT_PATH} + valid_anchor_manifest_path: ${oc.env:XTRAINER2_PACK_PHONE_MANIFEST_PATH} + temporal_proposal_path: ${oc.env:OPTR_TEMPORAL_PROPOSAL_PATH} + transform_rng_seed: 298226 + +trainer: + use_deepspeed: true + max_train_steps: 52380 + num_warmup_steps: 5000 + save_interval: null + save_steps: [26190, 34920, 43650, 52380] + eval_interval: null + repeated_diffusion_steps: 4 + pretrained_checkpoint: null + is_resume: true + allow_model_only_resume: false + resume_accelerate_checkpoint: ${oc.env:XTRAINER2_PACK_PHONE_RESUME_DIR} + branch_parent_seed: 42 + branch_rng_seed: 117465 + model_rng_seed: 588643 + stop_after_steps: 52380 + preflight_validate_first_batch: true + learning_rate: + base: 1.0e-05 + qwen_vl_interface: 1.0e-05 + action_model: 1.0e-04 + lr_scheduler_type: cosine_with_min_lr + scheduler_specific_kwargs: + min_lr: 5.0e-07 + freeze_modules: "" + trainable_modules: "" + frozen_modules_eval: false + logging_frequency: 50 + gradient_clipping: 1.0 + gradient_accumulation_steps: 4 + save_format: pt + save_final_checkpoint: false + optimizer: + name: AdamW + betas: [0.9, 0.95] + eps: 1.0e-08 + weight_decay: 1.0e-08 diff --git a/real_robot/source/configs/florence_xtrainer2_pack_phone_residual_resume_smoke.yaml b/real_robot/source/configs/florence_xtrainer2_pack_phone_residual_resume_smoke.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7020750dade524e6a5a867e4aadbd8638aec5883 --- /dev/null +++ b/real_robot/source/configs/florence_xtrainer2_pack_phone_residual_resume_smoke.yaml @@ -0,0 +1,104 @@ +# External exact-resume smoke. It uses the production batch geometry and +# common-start, but writes to an independent run and performs at most two +# optimizer updates. +run_id: florence-xtrainer2-pack-phone-residual-resume-smoke +run_root_dir: ${oc.env:OPTR_RUN_ROOT} +seed: 117465 +wandb_entity: null +wandb_project: OPTR-Florence-Xtrainer2 +wandb_group: pack-phone-optr +wandb_job_type: smoke +wandb_tags: [florence, xtrainer2, pack-phone, residual, resume-smoke] +is_debug: false +version_id: "0.21" + +framework: + name: QwenGR00T + qwenvl: + base_vlm: ${oc.env:FLORENCE_MODEL_PATH} + attn_implementation: eager + vl_hidden_dim: 1024 + selected_view_indices: [0, 1, 2] + compile_vision_tower: false + action_model: + action_model_type: DiT-B + hidden_size: 1024 + add_pos_embed: true + max_seq_len: 1024 + action_dim: 20 + state_dim: 20 + action_horizon: 16 + future_action_window_size: 15 + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + num_inference_timesteps: 4 + num_target_vision_tokens: 32 + diffusion_model_cfg: + cross_attention_dim: 1024 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + norm_type: ada_norm + num_layers: 16 + output_dim: 1024 + positional_embeddings: null + +datasets: + vla_data: + dataset_py: lerobot_datasets + data_root_dir: ${oc.env:XTRAINER2_DATA_ROOT} + data_mix: xtrainer2_pack_phone + action_type: abs_eef + action_mode: abs + include_state: true + per_device_batch_size: 4 + image_size: [224, 224] + video_backend: pyav + num_workers: 8 + expected_world_size: 8 + traceable_sampling: true + dataset_provenance_path: ${oc.env:XTRAINER2_PACK_PHONE_PROVENANCE_PATH} + dataset_content_fingerprint_path: ${oc.env:XTRAINER2_PACK_PHONE_FINGERPRINT_PATH} + valid_anchor_manifest_path: ${oc.env:XTRAINER2_PACK_PHONE_MANIFEST_PATH} + temporal_proposal_path: ${oc.env:OPTR_TEMPORAL_PROPOSAL_PATH} + transform_rng_seed: 298226 + +trainer: + use_deepspeed: true + max_train_steps: 52380 + num_warmup_steps: 5000 + save_interval: null + save_steps: [] + eval_interval: null + repeated_diffusion_steps: 4 + pretrained_checkpoint: null + is_resume: true + allow_model_only_resume: false + resume_accelerate_checkpoint: ${oc.env:XTRAINER2_PACK_PHONE_RESUME_DIR} + branch_parent_seed: 42 + branch_rng_seed: 117465 + model_rng_seed: 588643 + stop_after_steps: 17462 + preflight_validate_first_batch: true + learning_rate: + base: 1.0e-05 + qwen_vl_interface: 1.0e-05 + action_model: 1.0e-04 + lr_scheduler_type: cosine_with_min_lr + scheduler_specific_kwargs: + min_lr: 5.0e-07 + freeze_modules: "" + trainable_modules: "" + frozen_modules_eval: false + logging_frequency: 1 + gradient_clipping: 1.0 + gradient_accumulation_steps: 4 + save_format: pt + save_final_checkpoint: false + optimizer: + name: AdamW + betas: [0.9, 0.95] + eps: 1.0e-08 + weight_decay: 1.0e-08 diff --git a/real_robot/source/configs/florence_xtrainer2_pack_phone_smoke.yaml b/real_robot/source/configs/florence_xtrainer2_pack_phone_smoke.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e9762abc0acb52267dc606504a884e287386af3d --- /dev/null +++ b/real_robot/source/configs/florence_xtrainer2_pack_phone_smoke.yaml @@ -0,0 +1,117 @@ +# Smoke test for the pack_phone uniform run. Six optimizer steps with a +# checkpoint at step 4, which exercises the ZeRO-3 gather-and-consolidate save +# path that would otherwise only be reached after ~10 epochs. +# Keeps per_device_batch_size, accumulation, and world size identical to the +# real config so memory and the global batch of 128 are representative. +# vision-tower compile is off here to keep startup fast; the real run enables it. +run_id: florence-xtrainer2-pack-phone-smoke +run_root_dir: ${oc.env:OPTR_RUN_ROOT} +seed: 42 +wandb_entity: null +wandb_project: OPTR-Florence-Xtrainer2 +wandb_group: pack-phone +wandb_job_type: smoke +wandb_tags: + - florence + - xtrainer2 + - pack-phone + - smoke +is_debug: false +version_id: "0.21" + +framework: + name: QwenGR00T + qwenvl: + # NOTE: the backbone is selected by substring match on this path + # (starVLA/model/modules/vlm/__init__.py). The resolved value MUST contain + # "florence" or startup fails with NotImplementedError. + base_vlm: ${oc.env:FLORENCE_MODEL_PATH} + attn_implementation: eager + vl_hidden_dim: 1024 + selected_view_indices: [0, 1, 2] + compile_vision_tower: false + action_model: + action_model_type: DiT-B + hidden_size: 1024 + add_pos_embed: true + max_seq_len: 1024 + action_dim: 20 + state_dim: 20 + action_horizon: 16 + future_action_window_size: 15 + # NOTE: the effective flow-matching repeat factor is trainer.repeated_diffusion_steps. + # A framework-level key here would be silently ignored; do not re-add it. + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + num_inference_timesteps: 4 + num_target_vision_tokens: 32 + diffusion_model_cfg: + cross_attention_dim: 1024 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + norm_type: ada_norm + num_layers: 16 + output_dim: 1024 + positional_embeddings: null + +datasets: + vla_data: + dataset_py: lerobot_datasets + data_root_dir: ${oc.env:XTRAINER2_DATA_ROOT} + data_mix: xtrainer2_pack_phone + action_type: abs_eef + action_mode: abs + include_state: true + # 4 x 8 GPUs x accum 4 = global batch 128. Measured peaks on idle 80 GB + # cards: batch 4 = 60.3 GB, batch 5 = 73.4 GB, batch 6 OOMs. Holding + # global batch at 128 also requires per_device * accum == 16, so 4/4 is + # the largest usable microbatch; 8/2 does not fit in memory. + # See scripts/probe_batch_size.sh. + per_device_batch_size: 4 + image_size: [224, 224] + video_backend: pyav + num_workers: 2 + expected_world_size: 8 + traceable_sampling: true + dataset_provenance_path: ${oc.env:XTRAINER2_PACK_PHONE_PROVENANCE_PATH} + dataset_content_fingerprint_path: ${oc.env:XTRAINER2_PACK_PHONE_FINGERPRINT_PATH} + valid_anchor_manifest_path: null + temporal_proposal_path: null + +trainer: + use_deepspeed: true + max_train_steps: 6 + num_warmup_steps: 2 + save_interval: null + save_steps: [4] + eval_interval: null + # Effective flow-matching repeat factor (framework.action_model has no say). + # Per-DiT micro-batch is per_device_batch_size * this = 4 * 4 = 16. + repeated_diffusion_steps: 4 + learning_rate: + base: 1.0e-05 + qwen_vl_interface: 1.0e-05 + action_model: 1.0e-04 + lr_scheduler_type: cosine_with_min_lr + scheduler_specific_kwargs: + min_lr: 5.0e-07 + freeze_modules: "" + trainable_modules: "" + frozen_modules_eval: false + logging_frequency: 1 + gradient_clipping: 1.0 + # MUST equal deepspeed_config.gradient_accumulation_steps in the accelerate + # yaml. DeepSpeed's value governs the optimizer; this one sizes the dataloader + # truncation. A mismatch is only a warning, so keep them in sync by hand. + # Global batch = per_device 4 * 8 GPUs * accum 4 = 128. + gradient_accumulation_steps: 4 + save_format: pt + save_final_checkpoint: false + optimizer: + name: AdamW + betas: [0.9, 0.95] + eps: 1.0e-08 + weight_decay: 1.0e-08 diff --git a/real_robot/source/configs/florence_xtrainer2_pack_phone_uniform.yaml b/real_robot/source/configs/florence_xtrainer2_pack_phone_uniform.yaml new file mode 100644 index 0000000000000000000000000000000000000000..209fa4dfe51a54f342e022d26fa348175648b00e --- /dev/null +++ b/real_robot/source/configs/florence_xtrainer2_pack_phone_uniform.yaml @@ -0,0 +1,116 @@ +run_id: florence-xtrainer2-pack-phone-uniform +run_root_dir: ${oc.env:OPTR_RUN_ROOT} +seed: 42 +wandb_entity: null +wandb_project: OPTR-Florence-Xtrainer2 +wandb_group: pack-phone +wandb_job_type: train +wandb_tags: + - florence + - xtrainer2 + - pack-phone + - train +is_debug: false +version_id: "0.21" + +framework: + name: QwenGR00T + qwenvl: + # NOTE: the backbone is selected by substring match on this path + # (starVLA/model/modules/vlm/__init__.py). The resolved value MUST contain + # "florence" or startup fails with NotImplementedError. + base_vlm: ${oc.env:FLORENCE_MODEL_PATH} + attn_implementation: eager + vl_hidden_dim: 1024 + selected_view_indices: [0, 1, 2] + # Must stay false: torch.compile with mode="reduce-overhead" (CUDA graphs, + # Florence2.py:150-156) fails under ZeRO-3 with + # "AssertionError: wrong number of dimensions" from cudagraph_trees' + # assert_size_stride, because gathered parameter shapes differ between the + # trace and the replay. Verified failing on 8 ranks. + compile_vision_tower: false + action_model: + action_model_type: DiT-B + hidden_size: 1024 + add_pos_embed: true + max_seq_len: 1024 + action_dim: 20 + state_dim: 20 + action_horizon: 16 + future_action_window_size: 15 + # NOTE: the effective flow-matching repeat factor is trainer.repeated_diffusion_steps. + # A framework-level key here would be silently ignored; do not re-add it. + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + num_inference_timesteps: 4 + num_target_vision_tokens: 32 + diffusion_model_cfg: + cross_attention_dim: 1024 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + norm_type: ada_norm + num_layers: 16 + output_dim: 1024 + positional_embeddings: null + +datasets: + vla_data: + dataset_py: lerobot_datasets + data_root_dir: ${oc.env:XTRAINER2_DATA_ROOT} + data_mix: xtrainer2_pack_phone + action_type: abs_eef + action_mode: abs + include_state: true + # 4 x 8 GPUs x accum 4 = global batch 128. Measured peaks on idle 80 GB + # cards: batch 4 = 60.3 GB, batch 5 = 73.4 GB, batch 6 OOMs. Holding + # global batch at 128 also requires per_device * accum == 16, so 4/4 is + # the largest usable microbatch; 8/2 does not fit in memory. + # See scripts/probe_batch_size.sh. + per_device_batch_size: 4 + image_size: [224, 224] + video_backend: pyav + num_workers: 8 + expected_world_size: 8 + traceable_sampling: true + dataset_provenance_path: ${oc.env:XTRAINER2_PACK_PHONE_PROVENANCE_PATH} + dataset_content_fingerprint_path: ${oc.env:XTRAINER2_PACK_PHONE_FINGERPRINT_PATH} + valid_anchor_manifest_path: null + temporal_proposal_path: null + +trainer: + use_deepspeed: true + max_train_steps: 52380 + num_warmup_steps: 5000 + save_interval: null + save_steps: [8730, 17460, 26190, 34920, 43650, 52380] + eval_interval: null + # Effective flow-matching repeat factor (framework.action_model has no say). + # Per-DiT micro-batch is per_device_batch_size * this = 4 * 4 = 16. + repeated_diffusion_steps: 4 + learning_rate: + base: 1.0e-05 + qwen_vl_interface: 1.0e-05 + action_model: 1.0e-04 + lr_scheduler_type: cosine_with_min_lr + scheduler_specific_kwargs: + min_lr: 5.0e-07 + freeze_modules: "" + trainable_modules: "" + frozen_modules_eval: false + logging_frequency: 50 + gradient_clipping: 1.0 + # MUST equal deepspeed_config.gradient_accumulation_steps in the accelerate + # yaml. DeepSpeed's value governs the optimizer; this one sizes the dataloader + # truncation. A mismatch is only a warning, so keep them in sync by hand. + # Global batch = per_device 4 * 8 GPUs * accum 4 = 128. + gradient_accumulation_steps: 4 + save_format: pt + save_final_checkpoint: false + optimizer: + name: AdamW + betas: [0.9, 0.95] + eps: 1.0e-08 + weight_decay: 1.0e-08 diff --git a/real_robot/source/configs/florence_xtrainer2_pick_bottle_rel_aug_smoke.yaml b/real_robot/source/configs/florence_xtrainer2_pick_bottle_rel_aug_smoke.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6cceab51a59bacbfa0a864c3bb724da9ee8f21ea --- /dev/null +++ b/real_robot/source/configs/florence_xtrainer2_pick_bottle_rel_aug_smoke.yaml @@ -0,0 +1,119 @@ +# Smoke test for pick_bottle rel-EE + image-aug. Six optimizer steps with a +# checkpoint at step 4, which exercises relative-action statistics, random +# crop, and the ZeRO-3 gather-and-consolidate save path. +run_id: florence-xtrainer2-pick-bottle-rel-aug-smoke +run_root_dir: ${oc.env:OPTR_RUN_ROOT} +seed: 42 +wandb_entity: null +wandb_project: OPTR-Florence-Xtrainer2 +wandb_group: pick-bottle +wandb_job_type: smoke +wandb_tags: + - florence + - xtrainer2 + - pick-bottle + - rel-eef + - image-aug + - smoke +is_debug: false +version_id: "0.21" + +framework: + name: QwenGR00T + qwenvl: + # NOTE: the backbone is selected by substring match on this path + # (starVLA/model/modules/vlm/__init__.py). The resolved value MUST contain + # "florence" or startup fails with NotImplementedError. + base_vlm: ${oc.env:FLORENCE_MODEL_PATH} + attn_implementation: eager + vl_hidden_dim: 1024 + selected_view_indices: [0, 1, 2] + compile_vision_tower: false + action_model: + action_model_type: DiT-B + hidden_size: 1024 + add_pos_embed: true + max_seq_len: 1024 + action_dim: 20 + state_dim: 20 + action_horizon: 16 + future_action_window_size: 15 + # NOTE: the effective flow-matching repeat factor is trainer.repeated_diffusion_steps. + # A framework-level key here would be silently ignored; do not re-add it. + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + num_inference_timesteps: 4 + num_target_vision_tokens: 32 + diffusion_model_cfg: + cross_attention_dim: 1024 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + norm_type: ada_norm + num_layers: 16 + output_dim: 1024 + positional_embeddings: null + +datasets: + vla_data: + dataset_py: lerobot_datasets + data_root_dir: ${oc.env:XTRAINER2_DATA_ROOT} + data_mix: xtrainer2_pick_bottle + action_type: rel_eef + action_mode: rel + action_mode_apply_keys: + - action.left_position + - action.left_rotation + - action.right_position + - action.right_rotation + include_state: true + # 4 x 8 GPUs x accum 4 = global batch 128. Same memory envelope as the + # production run. + per_device_batch_size: 4 + image_size: [224, 224] + video_backend: pyav + num_workers: 2 + expected_world_size: 8 + traceable_sampling: true + transform_rng_seed: 298226 + dataset_provenance_path: ${oc.env:XTRAINER2_PICK_BOTTLE_PROVENANCE_PATH} + dataset_content_fingerprint_path: ${oc.env:XTRAINER2_PICK_BOTTLE_FINGERPRINT_PATH} + valid_anchor_manifest_path: null + temporal_proposal_path: null + +trainer: + use_deepspeed: true + max_train_steps: 6 + num_warmup_steps: 2 + save_interval: null + save_steps: [4] + eval_interval: null + # Effective flow-matching repeat factor (framework.action_model has no say). + # Per-DiT micro-batch is per_device_batch_size * this = 4 * 4 = 16. + repeated_diffusion_steps: 4 + learning_rate: + base: 1.0e-05 + qwen_vl_interface: 1.0e-05 + action_model: 1.0e-04 + lr_scheduler_type: cosine_with_min_lr + scheduler_specific_kwargs: + min_lr: 5.0e-07 + freeze_modules: "" + trainable_modules: "" + frozen_modules_eval: false + logging_frequency: 1 + gradient_clipping: 1.0 + # MUST equal deepspeed_config.gradient_accumulation_steps in the accelerate + # yaml. DeepSpeed's value governs the optimizer; this one sizes the dataloader + # truncation. A mismatch is only a warning, so keep them in sync by hand. + # Global batch = per_device 4 * 8 GPUs * accum 4 = 128. + gradient_accumulation_steps: 4 + save_format: pt + save_final_checkpoint: false + optimizer: + name: AdamW + betas: [0.9, 0.95] + eps: 1.0e-08 + weight_decay: 1.0e-08 diff --git a/real_robot/source/configs/florence_xtrainer2_pick_bottle_rel_aug_uniform.yaml b/real_robot/source/configs/florence_xtrainer2_pick_bottle_rel_aug_uniform.yaml new file mode 100644 index 0000000000000000000000000000000000000000..40971a93227a6a5b9e75c240716418ed94df15d0 --- /dev/null +++ b/real_robot/source/configs/florence_xtrainer2_pick_bottle_rel_aug_uniform.yaml @@ -0,0 +1,123 @@ +run_id: florence-xtrainer2-pick-bottle-rel-aug-uniform +run_root_dir: ${oc.env:OPTR_RUN_ROOT} +seed: 42 +wandb_entity: null +wandb_project: OPTR-Florence-Xtrainer2 +wandb_group: pick-bottle +wandb_job_type: train +wandb_tags: + - florence + - xtrainer2 + - pick-bottle + - rel-eef + - image-aug + - train +is_debug: false +version_id: "0.21" + +framework: + name: QwenGR00T + qwenvl: + # NOTE: the backbone is selected by substring match on this path + # (starVLA/model/modules/vlm/__init__.py). The resolved value MUST contain + # "florence" or startup fails with NotImplementedError. + base_vlm: ${oc.env:FLORENCE_MODEL_PATH} + attn_implementation: eager + vl_hidden_dim: 1024 + selected_view_indices: [0, 1, 2] + # Must stay false: torch.compile with mode="reduce-overhead" (CUDA graphs, + # Florence2.py:150-156) fails under ZeRO-3 with + # "AssertionError: wrong number of dimensions" from cudagraph_trees' + # assert_size_stride, because gathered parameter shapes differ between the + # trace and the replay. Verified failing on 8 ranks. + compile_vision_tower: false + action_model: + action_model_type: DiT-B + hidden_size: 1024 + add_pos_embed: true + max_seq_len: 1024 + action_dim: 20 + state_dim: 20 + action_horizon: 16 + future_action_window_size: 15 + # NOTE: the effective flow-matching repeat factor is trainer.repeated_diffusion_steps. + # A framework-level key here would be silently ignored; do not re-add it. + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + num_inference_timesteps: 4 + num_target_vision_tokens: 32 + diffusion_model_cfg: + cross_attention_dim: 1024 + dropout: 0.2 + final_dropout: true + interleave_self_attention: true + norm_type: ada_norm + num_layers: 16 + output_dim: 1024 + positional_embeddings: null + +datasets: + vla_data: + dataset_py: lerobot_datasets + data_root_dir: ${oc.env:XTRAINER2_DATA_ROOT} + data_mix: xtrainer2_pick_bottle + action_type: rel_eef + action_mode: rel + # Gripper stays absolute; xyz subtracts the current EE; quaternions use + # q_rel = inv(q_s0) ⊗ q_a rather than component-wise subtraction. + action_mode_apply_keys: + - action.left_position + - action.left_rotation + - action.right_position + - action.right_rotation + include_state: true + # 4 x 8 GPUs x accum 4 = global batch 128. Same memory envelope as pack_phone. + per_device_batch_size: 4 + image_size: [224, 224] + video_backend: pyav + num_workers: 8 + expected_world_size: 8 + traceable_sampling: true + transform_rng_seed: 298226 + dataset_provenance_path: ${oc.env:XTRAINER2_PICK_BOTTLE_PROVENANCE_PATH} + dataset_content_fingerprint_path: ${oc.env:XTRAINER2_PICK_BOTTLE_FINGERPRINT_PATH} + valid_anchor_manifest_path: null + temporal_proposal_path: null + +trainer: + use_deepspeed: true + # 75758 frames / global batch 128 ≈ 591.86 → 592 optimizer steps per data epoch. + max_train_steps: 35520 + num_warmup_steps: 3500 + save_interval: null + save_steps: [5920, 11840, 17760, 23680, 29600, 35520] + eval_interval: null + # Effective flow-matching repeat factor (framework.action_model has no say). + # Per-DiT micro-batch is per_device_batch_size * this = 4 * 4 = 16. + repeated_diffusion_steps: 4 + learning_rate: + base: 1.0e-05 + qwen_vl_interface: 1.0e-05 + action_model: 1.0e-04 + lr_scheduler_type: cosine_with_min_lr + scheduler_specific_kwargs: + min_lr: 5.0e-07 + freeze_modules: "" + trainable_modules: "" + frozen_modules_eval: false + logging_frequency: 50 + gradient_clipping: 1.0 + # MUST equal deepspeed_config.gradient_accumulation_steps in the accelerate + # yaml. DeepSpeed's value governs the optimizer; this one sizes the dataloader + # truncation. A mismatch is only a warning, so keep them in sync by hand. + # Global batch = per_device 4 * 8 GPUs * accum 4 = 128. + gradient_accumulation_steps: 4 + save_format: pt + save_final_checkpoint: false + optimizer: + name: AdamW + betas: [0.9, 0.95] + eps: 1.0e-08 + weight_decay: 1.0e-08 diff --git a/real_robot/source/deployment/__init__.py b/real_robot/source/deployment/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/real_robot/source/deployment/model_server/__init__.py b/real_robot/source/deployment/model_server/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/real_robot/source/deployment/model_server/policy_norm_processor.py b/real_robot/source/deployment/model_server/policy_norm_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..6103826c409bc8b15f95ab324d672a10c71b5501 --- /dev/null +++ b/real_robot/source/deployment/model_server/policy_norm_processor.py @@ -0,0 +1,538 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License. +"""PolicyNormProcessor — reuse the training-time ComposedModalityTransform. + +This class replaces the hand-rolled un-normalization math that previously +lived in every eval client. It rebuilds the *exact* ``ComposedModalityTransform`` +used at training time from a checkpoint: + + 1. Read ``config.yaml`` next to the ckpt → resolve ``data_mix`` → + look up ``robot_type`` from ``DATASET_NAMED_MIXTURES`` → + fetch the ``DataConfig`` from ``ROBOT_TYPE_CONFIG_MAP``. + 2. Build the transform pipeline via ``data_config.transform()``. + 3. Reconstruct a ``DatasetMetadata`` from the saved + ``dataset_statistics.json`` (which stores **combined** per-modality + arrays of length ``D``) by splitting it into per-key entries that match + ``data_config.action_keys`` / ``state_keys``. + 4. ``set_metadata(...)`` binds the metadata into every transform. + +After construction the caller simply invokes :meth:`unapply_actions` / +:meth:`apply_state` to get/normalize tensors using the same code path as +training — there is no second source of truth for normalization math. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional, Sequence + +import numpy as np +import torch + +from starVLA.dataloader.gr00t_lerobot.registry import ( + DATASET_NAMED_MIXTURES, + ROBOT_TYPE_CONFIG_MAP, +) +from starVLA.dataloader.gr00t_lerobot.schema import ( + DatasetMetadata, + StateActionMetadata, +) +from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform +from starVLA.model.framework.share_tools import read_mode_config + +logger = logging.getLogger(__name__) + + +def _resolve_robot_type( + model_cfg: dict, + unnorm_key: Optional[str] = None, +) -> str: + """Look up the training robot_type from the saved cfg. + + Convention used by starVLA train scripts: + ``cfg.datasets.vla_data.data_mix`` is a key in + ``DATASET_NAMED_MIXTURES`` whose value is a list of + ``(dataset_name, weight, robot_type)`` tuples. + + When a data_mix contains entries from multiple robot types (e.g. + ``bridge_rt_1`` covers ``oxe_bridge`` + ``oxe_rt1``), ``unnorm_key`` + is used to identify which embodiment is requested. In those mixtures + the ``robot_type`` field of each entry **matches** the top-level key in + ``dataset_statistics.json``, so ``unnorm_key`` serves as the selector. + """ + try: + data_mix = model_cfg["datasets"]["vla_data"]["data_mix"] + except (KeyError, TypeError) as e: + raise KeyError( + "ckpt config.yaml is missing `datasets.vla_data.data_mix`; " + "cannot resolve training-time robot_type." + ) from e + + if data_mix not in DATASET_NAMED_MIXTURES: + raise KeyError( + f"data_mix={data_mix!r} not in DATASET_NAMED_MIXTURES " + f"(available: {sorted(DATASET_NAMED_MIXTURES.keys())[:20]} ...). " + "Did you forget to register the example under examples//train_files/data_registry/?" + ) + + mixture = DATASET_NAMED_MIXTURES[data_mix] + robot_types = sorted({entry[2] for entry in mixture}) + + if len(robot_types) == 1: + return robot_types[0] + + # Multiple robot types in the mixture. + # Use unnorm_key as a direct selector: for multi-robot mixtures the + # dataset_statistics.json top-level keys equal the robot_type values. + if unnorm_key is not None and unnorm_key in robot_types: + return unnorm_key + + raise ValueError( + f"data_mix={data_mix!r} contains multiple robot_types {robot_types}. " + "Pass `unnorm_key` matching one of them to disambiguate " + f"(e.g. unnorm_key={robot_types[0]!r})." + ) + + +def _infer_key_dims( + data_config: Any, + combined_stats: Dict[str, Any], + modality_keys: Sequence[str], + modality: str, +) -> Dict[str, int]: + """Compute per-key dimensions for splitting combined stats arrays. + + Lookup priority: + 1. ``data_config._key_dims`` — explicit dict classvar on the DataConfig + (required for DataConfigs with non-uniform or multi-dim keys). + 2. Infer uniformly from stats array length: if ``D_total / n_keys`` is an + integer, use that as the uniform per-key dim. + 3. Fall back to dim=1 when no stats are available (empty combined dict). + + Args: + data_config: The DataConfig instance for the current robot type. + combined_stats: The raw stats dict for the chosen unnorm_key (the value + of ``norm_stats[unnorm_key]`` from ``dataset_statistics.json``). + modality_keys: Ordered list of full keys for this modality + (e.g. ``["action.left_joints", ...]``). + modality: ``"action"`` or ``"state"``. + + Returns: + Dict mapping each full key to its integer dimension. + """ + attr = f"{modality}_key_dims" + if hasattr(data_config, attr): + return dict(getattr(data_config, attr)) + + combined = combined_stats.get(modality, {}) + stat_arr = next((v for k, v in combined.items() if k != "mask"), None) + n_keys = len(modality_keys) + if stat_arr is not None and n_keys > 0: + D_total = len(stat_arr) + if D_total == n_keys: + return {k: 1 for k in modality_keys} + elif n_keys > 0 and D_total % n_keys == 0: + dim = D_total // n_keys + return {k: dim for k in modality_keys} + else: + raise ValueError( + f"Cannot infer per-key dims for modality={modality!r}: " + f"D_total={D_total} is not evenly divisible by n_keys={n_keys}. " + f"Add `{attr} = {{key: dim, ...}}` to the DataConfig " + f"(keys={list(modality_keys)})." + ) + return {k: 1 for k in modality_keys} + + +#: Width of each rotation representation, used to size a key after +#: ``StateActionTransform`` has converted it (see ``target_rotations``). +ROTATION_DIMS: Dict[str, int] = { + "axis_angle": 3, + "euler_angles": 3, + "quaternion": 4, + "rotation_6d": 6, + "matrix": 9, +} + + +def _rotation_dim(rotation_type: str) -> int: + """Width of ``rotation_type``; ``euler_angles_*`` carries a convention suffix.""" + if rotation_type.startswith("euler_angles"): + return ROTATION_DIMS["euler_angles"] + if rotation_type not in ROTATION_DIMS: + raise ValueError( + f"Unknown rotation representation {rotation_type!r}; " + f"expected one of {sorted(ROTATION_DIMS)}" + ) + return ROTATION_DIMS[rotation_type] + + +def _target_rotations(transform: ComposedModalityTransform) -> Dict[str, str]: + """Collect ``key -> target rotation`` from every StateActionTransform.""" + targets: Dict[str, str] = {} + for sub_transform in getattr(transform, "transforms", []): + for key, rotation in getattr(sub_transform, "target_rotations", {}).items(): + targets[key] = rotation + return targets + + +def _model_space_dims( + raw_key_dims: Dict[str, int], + target_rotations: Dict[str, str], +) -> Dict[str, int]: + """Per-key widths *after* the training transforms. + + The stats file and the robot speak raw widths (a quaternion is 4); the model + tensor is what ``target_rotations`` produced (rotation_6d is 6). For + xtrainer2 this is the 16 -> 20 expansion. + """ + return { + key: _rotation_dim(target_rotations[key]) if key in target_rotations else dim + for key, dim in raw_key_dims.items() + } + + +def _build_dataset_metadata( + stats_for_key: Dict[str, Any], + embodiment_tag: Any, + action_keys: Sequence[str], + state_keys: Sequence[str], + action_key_dims: Optional[Dict[str, int]] = None, + state_key_dims: Optional[Dict[str, int]] = None, + rotation_types: Optional[Dict[str, str]] = None, +) -> DatasetMetadata: + """Convert the *combined* stats arrays from ``dataset_statistics.json`` + back into a per-subkey :class:`DatasetMetadata` matching what the training + pipeline produced. + + The saved ``dataset_statistics.json`` stores stats as flat arrays of + length ``D = sum(per_key_dims)``. For each ``"action."`` key we + slice out ``dim_k`` elements starting at the current cursor and store them + under ``statistics.action. = {"min": [v0..v_{k-1}], ...}``. + + Args: + stats_for_key: ``norm_stats[unnorm_key]`` from ``dataset_statistics.json``. + embodiment_tag: Embodiment tag from the DataConfig. + action_keys: Ordered list of full action keys. + state_keys: Ordered list of full state keys. + action_key_dims: Per-key dimension dict (``{full_key: dim_k}``). + Defaults to dim=1 for every key. + state_key_dims: Per-key dimension dict for state keys. + Defaults to dim=1 for every key. + rotation_types: ``full_key -> source rotation representation`` for keys + the training transforms rotate. Without it those transforms cannot + be rebuilt, because the dataset's ``modality.json`` is not shipped + with the checkpoint. + """ + if action_key_dims is None: + action_key_dims = {k: 1 for k in action_keys} + if state_key_dims is None: + state_key_dims = {k: 1 for k in state_keys} + rotation_types = rotation_types or {} + + def _split_combined( + combined: Dict[str, Sequence[float]], + keys: Sequence[str], + key_dims: Dict[str, int], + ): + """Split combined arrays into per-subkey dicts using per-key dims. + + ``combined`` looks like ``{"min": [..D..], "max": [..D..], "mask": [..D..], ...}``. + ``keys`` is the ordered list of full keys. + ``key_dims`` maps each full key to its integer dimension. + Returns ``(stats_per_subkey, meta_per_subkey)``. + """ + stats_per_subkey: Dict[str, Dict[str, List[float]]] = {} + meta_per_subkey: Dict[str, StateActionMetadata] = {} + cursor = 0 + for full_key in keys: + subkey = full_key.split(".", 1)[1] + dim_k = key_dims.get(full_key, 1) + per_key: Dict[str, List[float]] = {} + for stat_name, arr in combined.items(): + if stat_name == "mask": + continue + end = cursor + dim_k + if end > len(arr): + # Saved combined array shorter than expected (truncated + # pad channels etc.). Skip this stat field silently. + continue + per_key[stat_name] = [float(v) for v in arr[cursor:end]] + stats_per_subkey[subkey] = per_key + meta_per_subkey[subkey] = StateActionMetadata( + absolute=True, + rotation_type=rotation_types.get(full_key), + shape=(dim_k,), + continuous=True, + ) + cursor += dim_k + return stats_per_subkey, meta_per_subkey + + action_combined = stats_for_key.get("action", {}) + state_combined = stats_for_key.get("state", {}) + + action_stats, action_meta = _split_combined(action_combined, action_keys, action_key_dims) + state_stats, state_meta = _split_combined(state_combined, state_keys, state_key_dims) + + # Pydantic accepts dict input with field validators + return DatasetMetadata.model_validate( + { + "statistics": { + "state": state_stats, + "action": action_stats, + }, + "modalities": { + "video": {}, + "state": state_meta, + "action": action_meta, + }, + "embodiment_tag": embodiment_tag.value + if hasattr(embodiment_tag, "value") + else embodiment_tag, + } + ) + + +class PolicyNormProcessor: + """Server-side normalization helper backed by training-time transforms. + + Construct once per checkpoint; call :meth:`unapply_actions` to convert + a normalized action chunk ``(T, D)`` back to env-space. + + Args: + ckpt_path: Path to the ``*.pt`` checkpoint (the loader looks for + ``config.yaml`` and ``dataset_statistics.json`` two dirs up). + unnorm_key: Which top-level key in ``dataset_statistics.json`` to use. + ``None`` → auto-pick the only key. + """ + + def __init__(self, ckpt_path: str, unnorm_key: Optional[str] = None) -> None: + self._ckpt_path = str(ckpt_path) + cfg, norm_stats = read_mode_config(self._ckpt_path) + self._model_cfg = cfg + self._norm_stats = norm_stats + + # 3-early) Pick the requested unnorm_key (or auto-select) BEFORE + # resolving robot_type so we can use it as a hint for multi-robot mixtures. + if unnorm_key is None: + if len(norm_stats) == 1: + unnorm_key = next(iter(norm_stats.keys())) + # else: defer error to step 3 below after robot_type resolution attempt + elif unnorm_key not in norm_stats: + raise KeyError( + f"unnorm_key={unnorm_key!r} not in {list(norm_stats.keys())}" + ) + self._unnorm_key = unnorm_key # may still be None for multi-key case + + # 1) Resolve which DataConfig was used at training. + robot_type = _resolve_robot_type(cfg, unnorm_key=unnorm_key) + if robot_type not in ROBOT_TYPE_CONFIG_MAP: + raise KeyError( + f"robot_type={robot_type!r} not in ROBOT_TYPE_CONFIG_MAP " + f"(available: {sorted(ROBOT_TYPE_CONFIG_MAP.keys())}). " + "Make sure the example's data_registry/data_config.py is importable." + ) + self._data_config = ROBOT_TYPE_CONFIG_MAP[robot_type] + self._action_keys: List[str] = list(self._data_config.action_keys) + self._state_keys: List[str] = list(getattr(self._data_config, "state_keys", [])) + self._modality_key_set = set(self._action_keys) | set(self._state_keys) + + # 2) Build training-time transform pipeline. + transform = self._data_config.transform() + if not isinstance(transform, ComposedModalityTransform): + transform = ComposedModalityTransform(transforms=[transform]) + self._transform = transform + + # 3) Pick the requested unnorm_key (finalize; error if still None here). + if self._unnorm_key is None: + raise ValueError( + f"Multiple unnorm_keys in dataset_statistics.json: " + f"{list(norm_stats.keys())}. Pass unnorm_key explicitly." + ) + unnorm_key = self._unnorm_key + + # 4) Resolve per-key dims (handles multi-d action/state keys). + stats_for_unnorm = norm_stats[unnorm_key] + self._action_key_dims: Dict[str, int] = _infer_key_dims( + self._data_config, stats_for_unnorm, self._action_keys, "action" + ) + self._state_key_dims: Dict[str, int] = _infer_key_dims( + self._data_config, stats_for_unnorm, self._state_keys, "state" + ) + + # 5) Widths of the model tensor, which differ from the raw widths for + # every key the transforms rotate (xtrainer2: quaternion 4 -> 6D). + target_rotations = _target_rotations(self._transform) + self._action_model_dims = _model_space_dims(self._action_key_dims, target_rotations) + self._state_model_dims = _model_space_dims(self._state_key_dims, target_rotations) + + # 6) Build & bind metadata. + ds_meta = _build_dataset_metadata( + stats_for_key=stats_for_unnorm, + embodiment_tag=self._data_config.embodiment_tag, + action_keys=self._action_keys, + state_keys=self._state_keys, + action_key_dims=self._action_key_dims, + state_key_dims=self._state_key_dims, + rotation_types=dict(getattr(self._data_config, "rotation_types", {})), + ) + # Only the state/action transforms are bound. The video transforms would + # demand per-camera resolutions that the checkpoint does not carry, and + # the server never runs them: images are resized by the framework's own + # inference path. + self._state_action_transforms = [ + sub_transform + for sub_transform in self._transform.transforms + if self._modality_key_set.intersection(getattr(sub_transform, "apply_to", [])) + ] + for sub_transform in self._state_action_transforms: + sub_transform.set_metadata(ds_meta) + sub_transform.eval() # mark transforms as eval-mode + + logger.info( + "PolicyNormProcessor ready: robot_type=%s, unnorm_key=%s, " + "action_keys=%s (raw dims=%s, model dims=%s), state_keys=%s " + "(raw dims=%s, model dims=%s)", + robot_type, + unnorm_key, + self._action_keys, + [self._action_key_dims[k] for k in self._action_keys], + [self._action_model_dims[k] for k in self._action_keys], + self._state_keys, + [self._state_key_dims[k] for k in self._state_keys], + [self._state_model_dims[k] for k in self._state_keys], + ) + + # ------------------------------------------------------------------ + # Public properties + # ------------------------------------------------------------------ + @property + def action_keys(self) -> List[str]: + return list(self._action_keys) + + @property + def state_keys(self) -> List[str]: + return list(self._state_keys) + + @property + def unnorm_key(self) -> str: + return self._unnorm_key + + @property + def available_unnorm_keys(self) -> List[str]: + return list(self._norm_stats.keys()) + + @property + def transform(self) -> ComposedModalityTransform: + return self._transform + + @property + def action_dim(self) -> int: + """Width of the model action vector (post-transform).""" + return sum(self._action_model_dims[key] for key in self._action_keys) + + @property + def raw_action_dim(self) -> int: + """Width of an action in robot/dataset coordinates.""" + return sum(self._action_key_dims[key] for key in self._action_keys) + + @property + def state_dim(self) -> int: + """Width of the model state vector (post-transform).""" + return sum(self._state_model_dims[key] for key in self._state_keys) + + @property + def raw_state_dim(self) -> int: + """Width of a state as the robot reports it.""" + return sum(self._state_key_dims[key] for key in self._state_keys) + + # ------------------------------------------------------------------ + # Forward path (robot observation → model state) + # ------------------------------------------------------------------ + def apply_state(self, raw_state: np.ndarray) -> np.ndarray: + """Normalize a raw state exactly as the training dataloader did. + + Args: + raw_state: ``(raw_state_dim,)`` or ``(T, raw_state_dim)`` in dataset + coordinates -- i.e. after any robot-specific reordering, with + rotations still in their source representation. + + Returns: + ``(T, state_dim)`` normalized state, the layout the action head was + trained on. + """ + state = np.asarray(raw_state, dtype=np.float32) + if state.ndim == 1: + state = state[None, :] + if state.ndim != 2: + raise ValueError(f"Expected (raw_state_dim,) or (T, raw_state_dim); got {state.shape}") + if state.shape[-1] != self.raw_state_dim: + raise ValueError( + f"Expected {self.raw_state_dim} raw state dims " + f"(state_keys={self._state_keys}, dims=" + f"{[self._state_key_dims[k] for k in self._state_keys]}); got {state.shape[-1]}" + ) + + data: Dict[str, Any] = {} + cursor = 0 + for full_key in self._state_keys: + dim_k = self._state_key_dims[full_key] + data[full_key] = state[:, cursor : cursor + dim_k] + cursor += dim_k + + for sub_transform in self._state_action_transforms: + data = sub_transform(data) + + parts: List[np.ndarray] = [] + for full_key in self._state_keys: + value = data[full_key] + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + parts.append(np.asarray(value, dtype=np.float32)) + return np.concatenate(parts, axis=-1) + + # ------------------------------------------------------------------ + # Inverse path (model output → env action) + # ------------------------------------------------------------------ + def unapply_actions(self, normalized_actions: np.ndarray) -> np.ndarray: + """Invert action normalization using the training-time pipeline. + + Args: + normalized_actions: shape ``(T, D)`` where ``D == action_dim``. + + Returns: + ``(T, raw_action_dim)`` un-normalized actions in env coordinates. + """ + normalized_actions = np.asarray(normalized_actions) + assert normalized_actions.ndim == 2, ( + f"Expected (T, D); got shape {normalized_actions.shape}" + ) + + # Split (T, D) into per-key {full_key: torch.Tensor[T, dim_k]}. + data: Dict[str, torch.Tensor] = {} + cursor = 0 + for full_key in self._action_keys: + dim_k = self._action_model_dims.get(full_key, 1) + slice_ = normalized_actions[..., cursor : cursor + dim_k] + data[full_key] = torch.as_tensor(slice_, dtype=torch.float32) + cursor += dim_k + + if cursor != normalized_actions.shape[-1]: + raise ValueError( + f"Sum of per-key model dims ({cursor}) != action_dim " + f"({normalized_actions.shape[-1]}). " + f"action_keys={self._action_keys}, " + f"action_model_dims={self._action_model_dims}" + ) + + out: Dict[str, Any] = data + for sub_transform in reversed(self._state_action_transforms): + out = sub_transform.unapply(out) + + parts: List[np.ndarray] = [] + for full_key in self._action_keys: + v = out[full_key] + if isinstance(v, torch.Tensor): + v = v.detach().cpu().numpy() + parts.append(np.asarray(v)) + return np.concatenate(parts, axis=-1) diff --git a/real_robot/source/deployment/model_server/policy_wrapper.py b/real_robot/source/deployment/model_server/policy_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..8898ca06144b0eb11a0be92764faac114ef4414d --- /dev/null +++ b/real_robot/source/deployment/model_server/policy_wrapper.py @@ -0,0 +1,280 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License. +"""Policy server wrapper. + +Encapsulates a `baseframework` instance plus a :class:`PolicyNormProcessor` +that reuses the *training-time* :class:`ComposedModalityTransform` for action +un-normalization (no hand-rolled math). The websocket server returns +already-unnormalized actions. + +Client-side responsibilities that REMAIN on the client: + - environment-specific adapters (image_history, gripper sticky, action + ensembling) + - chunk-cache scheduling (`step % chunk_size == 0` triggers a new infer) + +Embodiments whose DataConfig declares a ``robot_command_layout`` (currently +xtrainer2) additionally get the raw-robot conversions here, so the client never +has to reproduce training-time normalization or the quaternion order: + + - ``state`` arrives in robot coordinates (16-D, xyzw) and is reordered to the + training layout, then normalized with the training transforms (20-D) + - returned actions are un-normalized *and* reordered back to xyzw + +Exposed API: + - ``metadata`` (dict, sent at handshake): ``action_chunk_size``, + ``available_unnorm_keys``, ``action_keys``, ``state_keys``, plus the + embodiment's raw-robot contract. + - ``predict_action(examples, unnorm_key=None, **kwargs)`` returns + ``{"actions": np.ndarray[B, T, action_dim]}``. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +import numpy as np +import torch + +from starVLA.model.framework.base_framework import baseframework +from starVLA.model.framework.share_tools import read_mode_config + +from deployment.model_server.policy_norm_processor import PolicyNormProcessor +from deployment.model_server import xtrainer2_command + + +class PolicyServerWrapper: + """Wraps a `baseframework` for use as a websocket-server policy.""" + + def __init__( + self, + ckpt_path: str, + device: str = "cuda", + use_bf16: bool = False, + unnorm_key: Optional[str] = None, + normalized_action_clip: Optional[float] = None, + ) -> None: + self._ckpt_path = str(ckpt_path) + self._normalized_action_clip = normalized_action_clip + + logging.info("PolicyServerWrapper: loading framework from %s", self._ckpt_path) + framework = baseframework.from_pretrained(self._ckpt_path) + if use_bf16: + framework = framework.to(torch.bfloat16) + framework = framework.to(device).eval() + self._framework = framework + + # Co-located metadata. + model_cfg, _ = read_mode_config(self._ckpt_path) + self._model_cfg = model_cfg + vla_data_cfg = model_cfg.get("datasets", {}).get("vla_data", {}) + include_state = vla_data_cfg.get("include_state", False) + self._uses_state = include_state not in (False, None, 0, "False", "false", "0") + + # action_chunk_size = future_action_window_size + 1 (matches old client). + action_model_cfg = model_cfg["framework"]["action_model"] + + if "action_horizon" in action_model_cfg: + self._action_chunk_size = int(action_model_cfg["action_horizon"]) + elif "future_action_window_size" in action_model_cfg: + self._action_chunk_size = int(action_model_cfg["future_action_window_size"]) + 1 + else: + raise ValueError( + f"PolicyServerWrapper: no action_horizon or future_action_window_size found in model config for {self._ckpt_path}" + ) + # Cache of PolicyNormProcessor instances per unnorm_key. + # For single-dataset ckpts unnorm_key is auto-selected; for multi-dataset + # ckpts clients must pass unnorm_key per request. + self._default_unnorm_key = unnorm_key + self._norm_processors: Dict[str, PolicyNormProcessor] = {} + + # Peek at available keys without building a full processor. + _, _ns = read_mode_config(self._ckpt_path) + self._available_unnorm_keys: List[str] = list(_ns.keys()) + + # Eagerly build when unambiguous; defer for multi-key / no explicit key. + if unnorm_key is not None or len(self._available_unnorm_keys) == 1: + default_proc = self._get_processor(unnorm_key) + self._default_unnorm_key = default_proc.unnorm_key + logging.info( + "PolicyServerWrapper ready: action_chunk_size=%d, default_unnorm_key=%s, " + "available_unnorm_keys=%s, action_keys=%s, state_keys=%s", + self._action_chunk_size, + default_proc.unnorm_key, + default_proc.available_unnorm_keys, + default_proc.action_keys, + default_proc.state_keys, + ) + else: + logging.info( + "PolicyServerWrapper ready (multi-key): action_chunk_size=%d, " + "available_unnorm_keys=%s — clients must pass unnorm_key per request.", + self._action_chunk_size, + self._available_unnorm_keys, + ) + + def _get_processor(self, unnorm_key: Optional[str]) -> PolicyNormProcessor: + cache_key = unnorm_key if unnorm_key is not None else "__default__" + if cache_key not in self._norm_processors: + self._norm_processors[cache_key] = PolicyNormProcessor( + self._ckpt_path, unnorm_key=unnorm_key + ) + return self._norm_processors[cache_key] + + @staticmethod + def _uses_xtrainer2_layout(proc: PolicyNormProcessor) -> bool: + """Whether this embodiment speaks the raw xtrainer2 robot layout.""" + layout = getattr(proc._data_config, "robot_command_layout", None) + return layout == xtrainer2_command.COMMAND_LAYOUT + + @property + def metadata(self) -> Dict[str, Any]: + """Model-invariant metadata; sent to client at websocket handshake.""" + base = { + "env": "starvla_policy_server", + "ckpt_path": self._ckpt_path, + "action_chunk_size": self._action_chunk_size, + "uses_state": self._uses_state, + "available_unnorm_keys": self._available_unnorm_keys, + "default_unnorm_key": self._default_unnorm_key, + "normalized_action_clip": self._normalized_action_clip, + "action_head_attention_mask": False, + "inference_rng": "continuous_server_stream", + # Every client sends images this way; the framework does its own + # resize, so only the layout and channel order are binding. + "image_format": "HWC_RGB_uint8", + "image_key": "image", + } + # Enrich with per-embodiment keys when a default processor already exists. + if self._default_unnorm_key is not None: + proc = self._get_processor(self._default_unnorm_key) + base["action_keys"] = proc.action_keys + base["state_keys"] = proc.state_keys + base["camera_order"] = list(getattr(proc._data_config, "video_keys", [])) + base["model_action_dim"] = proc.action_dim + base["model_state_dim"] = proc.state_dim + base["server_normalizes_state"] = self._server_normalizes_state(proc) + if base["server_normalizes_state"]: + base["state_dim"] = proc.raw_state_dim + if self._uses_xtrainer2_layout(proc): + base["robot_command_layout"] = xtrainer2_command.COMMAND_LAYOUT + base["action_dim"] = proc.raw_action_dim + base["model_quaternion_format"] = xtrainer2_command.MODEL_QUATERNION_FORMAT + base["robot_quaternion_format"] = xtrainer2_command.SOURCE_QUATERNION_FORMAT + base["quaternion_starts"] = list(xtrainer2_command.QUATERNION_STARTS) + action_mode = str( + getattr(self, "_model_cfg", {}) + .get("datasets", {}) + .get("vla_data", {}) + .get("action_mode", "abs") + ) + base["action_mode"] = action_mode + base["server_composes_relative_eef"] = ( + xtrainer2_command.is_relative_action_mode(action_mode) + ) + return base + + @staticmethod + def _server_normalizes_state(proc: PolicyNormProcessor) -> bool: + """Whether ``state`` arrives raw and must be normalized server-side.""" + return bool(getattr(proc._data_config, "server_normalizes_state", False)) + + @classmethod + def _prepare_examples(cls, examples: List[dict], proc: PolicyNormProcessor) -> List[dict]: + """Turn raw-robot observations into the tensors the model was trained on. + + Leaves examples untouched for embodiments that do not opt in, so clients + that already send normalized state keep working. + """ + if not cls._server_normalizes_state(proc): + return examples + + to_model_layout = ( + xtrainer2_command.robot_state_to_model + if cls._uses_xtrainer2_layout(proc) + else None + ) + + prepared: List[dict] = [] + for example in examples: + if "state" not in example or example["state"] is None: + prepared.append(example) + continue + raw_state = np.asarray(example["state"], dtype=np.float32) + if raw_state.shape[-1] != proc.raw_state_dim: + raise ValueError( + f"state must be the robot's raw {proc.raw_state_dim}-D vector " + f"(state_keys={proc.state_keys}); got shape {raw_state.shape}. " + "The server applies the training-time normalization itself." + ) + if to_model_layout is not None: + raw_state = to_model_layout(raw_state) + example = dict(example) + # Keep the training-layout (wxyz) pose so rel-mode actions can be + # composed back to absolute commands after un-normalization. + example["_raw_state_wxyz"] = raw_state + example["state"] = proc.apply_state(raw_state) + prepared.append(example) + return prepared + + def predict_action( + self, + examples: List[dict], + unnorm_key: Optional[str] = None, + **kwargs, + ) -> Dict[str, np.ndarray]: + """Run the framework, then un-normalize via training-time transforms. + + Args: + examples: list of dicts (each with ``image`` / ``lang`` / optional + ``state``). For embodiments that declare + ``server_normalizes_state``, ``state`` is the robot's raw vector + and is normalized here. + unnorm_key: dataset key for un-normalization stats. ``None`` --> + use the wrapper's default (auto-picked at startup). + **kwargs: forwarded to the framework's ``predict_action`` + (``do_sample``, ``use_ddim``, ``num_ddim_steps``, ...). + + Returns: + ``{"actions": np.ndarray[B, T, D]}`` -- un-normalized. + """ + effective_key = unnorm_key if unnorm_key is not None else self._default_unnorm_key + if effective_key is None: + if len(self._available_unnorm_keys) == 1: + effective_key = self._available_unnorm_keys[0] + else: + raise ValueError( + f"predict_action: unnorm_key not specified and no default set. " + f"Pass one of {self._available_unnorm_keys}." + ) + proc = self._get_processor(effective_key) + examples = self._prepare_examples(examples, proc) + + out = self._framework.predict_action(examples=examples, **kwargs) + normalized = np.asarray(out["normalized_actions"]) # (B, T, D) + if self._normalized_action_clip is not None: + limit = float(self._normalized_action_clip) + normalized = np.clip(normalized, -limit, limit) + + unnorm = np.stack( + [proc.unapply_actions(normalized[b]) for b in range(normalized.shape[0])], + axis=0, + ) + if self._uses_xtrainer2_layout(proc): + if xtrainer2_command.is_relative_action_mode( + self._model_cfg.get("datasets", {}).get("vla_data", {}).get("action_mode") + ): + composed = [] + for batch_index in range(unnorm.shape[0]): + state_wxyz = examples[batch_index].get("_raw_state_wxyz") + if state_wxyz is None: + raise ValueError( + "action_mode=rel requires a raw robot state so relative " + "EE commands can be composed back to absolute poses" + ) + composed.append( + xtrainer2_command.compose_relative_eef_wxyz(unnorm[batch_index], state_wxyz) + ) + unnorm = np.stack(composed, axis=0) + unnorm = xtrainer2_command.model_actions_to_robot(unnorm) + return {"actions": unnorm} diff --git a/real_robot/source/deployment/model_server/server_policy.py b/real_robot/source/deployment/model_server/server_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..03b75dc083043716c89969b27acaa29e0e24d90a --- /dev/null +++ b/real_robot/source/deployment/model_server/server_policy.py @@ -0,0 +1,79 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); + +import logging +import socket +import argparse +from deployment.model_server.tools.websocket_policy_server import WebsocketPolicyServer +from deployment.model_server.policy_wrapper import PolicyServerWrapper +import os + + +def main(args) -> None: + # Example usage: + # policy = YourPolicyClass() # Replace with your actual policy class + # server = WebsocketPolicyServer(policy, host="localhost", port=10091) + # server.serve_forever() + + policy = PolicyServerWrapper( + ckpt_path=args.ckpt_path, + device="cuda", + use_bf16=args.use_bf16, + unnorm_key=args.unnorm_key, + ) + + hostname = socket.gethostname() + local_ip = socket.gethostbyname(hostname) + logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip) + # The handshake metadata is the request contract; log it so a client + # mismatch (camera order, state width, quaternion order) is visible here. + for key, value in sorted(policy.metadata.items()): + logging.info(" metadata %-24s %s", key, value) + + # start websocket server + server = WebsocketPolicyServer( + policy=policy, + host="0.0.0.0", + port=args.port, + idle_timeout=args.idle_timeout, + metadata={**policy.metadata, "policy_seed": args.seed}, + ) + logging.info("server running ...") + server.serve_forever() + + +def build_argparser(): + parser = argparse.ArgumentParser() + parser.add_argument("--ckpt_path", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct") + parser.add_argument("--port", type=int, default=10093) + parser.add_argument("--use_bf16", action="store_true") + parser.add_argument( + "--unnorm_key", + type=str, + default=None, + help="dataset_statistics.json key for un-normalization; required only for multi-dataset checkpoints", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--idle_timeout" , type=int, default=1800, help="Idle timeout in seconds, -1 means never close") + return parser + + +def start_debugpy_once(): + """start debugpy once""" + import debugpy + if getattr(start_debugpy_once, "_started", False): + return + debugpy.listen(("0.0.0.0", 10095)) + print("Waiting for VSCode attach on 0.0.0.0:10095 ...") + debugpy.wait_for_client() + start_debugpy_once._started = True + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + parser = build_argparser() + args = parser.parse_args() + if os.getenv("DEBUG", "").strip().lower() in {"1", "true", "yes"}: + print("DEBUGPY is enabled") + start_debugpy_once() + main(args) diff --git a/real_robot/source/deployment/model_server/tools/__init__.py b/real_robot/source/deployment/model_server/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/real_robot/source/deployment/model_server/tools/image_tools.py b/real_robot/source/deployment/model_server/tools/image_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..ba95f3bdb97a8e8f0e04f3f4e1cc78a461dc45af --- /dev/null +++ b/real_robot/source/deployment/model_server/tools/image_tools.py @@ -0,0 +1,123 @@ +import numpy as np +from PIL import Image + + +def convert_to_uint8(img: np.ndarray) -> np.ndarray: + """Converts an image to uint8 if it is a float image. + + This is important for reducing the size of the image when sending it over the network. + """ + if np.issubdtype(img.dtype, np.floating): + img = (255 * img).astype(np.uint8) + return img + + +def resize_with_pad(images: np.ndarray, height: int, width: int, method=Image.BILINEAR) -> np.ndarray: + """Replicates tf.image.resize_with_pad for multiple images using PIL. Resizes a batch of images to a target height. + + Args: + images: A batch of images in [..., height, width, channel] format. + height: The target height of the image. + width: The target width of the image. + method: The interpolation method to use. Default is bilinear. + + Returns: + The resized images in [..., height, width, channel]. + """ + # If the images are already the correct size, return them as is. + if images.shape[-3:-1] == (height, width): + return images + + original_shape = images.shape + + images = images.reshape(-1, *original_shape[-3:]) + resized = np.stack([_resize_with_pad_pil(Image.fromarray(im), height, width, method=method) for im in images]) + return resized.reshape(*original_shape[:-3], *resized.shape[-3:]) + + +def _resize_with_pad_pil(image: Image.Image, height: int, width: int, method: int) -> Image.Image: + """Replicates tf.image.resize_with_pad for one image using PIL. Resizes an image to a target height and + width without distortion by padding with zeros. + + Unlike the jax version, note that PIL uses [width, height, channel] ordering instead of [batch, h, w, c]. + """ + cur_width, cur_height = image.size + if cur_width == width and cur_height == height: + return image # No need to resize if the image is already the correct size. + + ratio = max(cur_width / width, cur_height / height) + resized_height = int(cur_height / ratio) + resized_width = int(cur_width / ratio) + resized_image = image.resize((resized_width, resized_height), resample=method) + + zero_image = Image.new(resized_image.mode, (width, height), 0) + pad_height = max(0, int((height - resized_height) / 2)) + pad_width = max(0, int((width - resized_width) / 2)) + zero_image.paste(resized_image, (pad_width, pad_height)) + assert zero_image.size == (width, height) + return zero_image + + +from typing import Any + + +def to_pil_preserve(images: Any, scale_float: bool = True): + """ + Convert (possibly nested) numpy image arrays back to PIL.Image WITHOUT changing spatial shape + or nesting structure. + + Accepts: + - np.ndarray with shape (H, W, C), C in {1,3,4}, dtype uint8 or float + - PIL.Image.Image (returned as-is) + - Nested list / tuple structures containing the above + + Guarantees: + - No resize / pad / crop performed + - Returns an object with the SAME nesting layout (list -> list, tuple -> tuple) + - Only dtype (float -> uint8) and channel-mode adaptation may happen + * float arrays assumed in [0,1] if scale_float=True (scaled *255 + clip) + Args: + images: input object / sequence + scale_float: whether to scale float images in [0,1] to uint8 + Returns: + Mirrored structure with all leaf nodes as PIL.Image.Image + """ + + def _convert(obj): + # Nested containers + if isinstance(obj, list): + return [_convert(x) for x in obj] + if isinstance(obj, tuple): + return tuple(_convert(x) for x in obj) + + # PIL stays + if isinstance(obj, Image.Image): + return obj + + # numpy -> PIL + if isinstance(obj, np.ndarray): + arr = obj + if arr.ndim != 3: + raise ValueError(f"Expected 3D array (H,W,C), got shape={arr.shape}") + if arr.shape[2] not in (1, 3, 4): + raise ValueError(f"Channel count must be 1/3/4, got {arr.shape[2]}") + if np.issubdtype(arr.dtype, np.floating): + if scale_float: + arr = np.clip(arr, 0.0, 1.0) + arr = (arr * 255.0 + 0.5).astype(np.uint8) + else: + raise TypeError("Float array provided but scale_float=False") + elif arr.dtype != np.uint8: + arr = arr.astype(np.uint8) + + # Single channel -> 'L' + if arr.shape[2] == 1: + arr = arr[:, :, 0] + return Image.fromarray(arr, mode="L") + # 3 channels -> RGB, 4 -> RGBA + mode = "RGB" if arr.shape[2] == 3 else "RGBA" + return Image.fromarray(arr, mode=mode) + + raise TypeError(f"Unsupported element type: {type(obj)}") + + return _convert(images) diff --git a/real_robot/source/deployment/model_server/tools/msgpack_numpy.py b/real_robot/source/deployment/model_server/tools/msgpack_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..d98b69a31325b9b4393a09001c8135002fea80f5 --- /dev/null +++ b/real_robot/source/deployment/model_server/tools/msgpack_numpy.py @@ -0,0 +1,59 @@ +"""Adds NumPy array support to msgpack. + +msgpack is good for (de)serializing data over a network for multiple reasons: +- msgpack is secure (as opposed to pickle/dill/etc which allow for arbitrary code execution) +- msgpack is widely used and has good cross-language support +- msgpack does not require a schema (as opposed to protobuf/flatbuffers/etc) which is convenient in dynamically typed + languages like Python and JavaScript +- msgpack is fast and efficient (as opposed to readable formats like JSON/YAML/etc); I found that msgpack was ~4x faster + than pickle for serializing large arrays using the below strategy + +The code below is adapted from https://github.com/lebedov/msgpack-numpy. +Its original BSD license is retained in licenses/msgpack-numpy-LICENSE. +The reason not to use that library directly is +that it falls back to pickle for object arrays. +""" + +import functools + +import msgpack +import numpy as np + + +def pack_array(obj): + if (isinstance(obj, (np.ndarray, np.generic))) and obj.dtype.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {obj.dtype}") + + if isinstance(obj, np.ndarray): + return { + b"__ndarray__": True, + b"data": obj.tobytes(), + b"dtype": obj.dtype.str, + b"shape": obj.shape, + } + + if isinstance(obj, np.generic): + return { + b"__npgeneric__": True, + b"data": obj.item(), + b"dtype": obj.dtype.str, + } + + return obj + + +def unpack_array(obj): + if b"__ndarray__" in obj: + return np.ndarray(buffer=obj[b"data"], dtype=np.dtype(obj[b"dtype"]), shape=obj[b"shape"]) + + if b"__npgeneric__" in obj: + return np.dtype(obj[b"dtype"]).type(obj[b"data"]) + + return obj + + +Packer = functools.partial(msgpack.Packer, default=pack_array) +packb = functools.partial(msgpack.packb, default=pack_array) + +Unpacker = functools.partial(msgpack.Unpacker, object_hook=unpack_array) +unpackb = functools.partial(msgpack.unpackb, object_hook=unpack_array) diff --git a/real_robot/source/deployment/model_server/tools/websocket_policy_client.py b/real_robot/source/deployment/model_server/tools/websocket_policy_client.py new file mode 100644 index 0000000000000000000000000000000000000000..c7c7a1fff814dc68ea0fc0aa1365330e92fbe01c --- /dev/null +++ b/real_robot/source/deployment/model_server/tools/websocket_policy_client.py @@ -0,0 +1,75 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); + +import logging, argparse +import time, os +from typing import Dict, Optional, Tuple + +from typing_extensions import override +import websockets.sync.client + +from . import msgpack_numpy + + +class WebsocketClientPolicy: + """Implements the Policy interface by communicating with a server over websocket. + + See WebsocketPolicyServer for a corresponding server implementation. + """ + + def __init__(self, host: str = "127.0.0.1", port: Optional[int] = 10093, api_key: Optional[str] = None) -> None: + # 0.0.0.0 cannot be used as a connection target, here default 127.0.0.1 + self._uri = f"ws://{host}" + if port is not None: + self._uri += f":{port}" + self._packer = msgpack_numpy.Packer() + self._api_key = api_key + self._ws, self._server_metadata = self._wait_for_server() + + def get_server_metadata(self) -> Dict: + return self._server_metadata + + def _wait_for_server(self, timeout: float = 300) -> Tuple[websockets.sync.client.ClientConnection, Dict]: + logging.info(f"Waiting for server at {self._uri}...") + start_time = time.time() + + for k in ("HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy"): + os.environ.pop(k, None) + + while True: + if time.time() - start_time > timeout: + raise TimeoutError(f"Failed to connect to server within {timeout} seconds") + + try: + headers = {"Authorization": f"Api-Key {self._api_key}"} if self._api_key else None + conn = websockets.sync.client.connect( + self._uri, + compression=None, + max_size=None, + additional_headers=headers, + open_timeout=150, + ping_interval=None, + ping_timeout=None, + ) + metadata = msgpack_numpy.unpackb(conn.recv()) + return conn, metadata + except ConnectionRefusedError: + logging.info(f"Still waiting for server {self._uri} ...") + time.sleep(2) + + def close(self) -> None: + try: + self._ws.close() + except Exception: + pass + + @override + def predict_action(self, query_info: Dict) -> Dict: + data = self._packer.pack(query_info) + self._ws.send(data) + response = self._ws.recv() + if isinstance(response, str): + raise RuntimeError(f"Error in inference server:\n{response}") + return msgpack_numpy.unpackb(response) + + diff --git a/real_robot/source/deployment/model_server/tools/websocket_policy_server.py b/real_robot/source/deployment/model_server/tools/websocket_policy_server.py new file mode 100644 index 0000000000000000000000000000000000000000..be9c9c6ce553f76606d2f1d99fa76831c285c576 --- /dev/null +++ b/real_robot/source/deployment/model_server/tools/websocket_policy_server.py @@ -0,0 +1,160 @@ +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); + +import asyncio +import logging +import traceback +import time + +import websockets.asyncio.server +import websockets.frames + +# from openpi_client import base_policy as _base_policy +from . import msgpack_numpy + +class WebsocketPolicyServer: + """Serves a policy using the websocket protocol. See websocket_client_policy.py for a client implementation. + + Currently only implements the `load` and `infer` methods. + """ + + def __init__( + self, + policy, + host: str = "0.0.0.0", + port: int = 10093, + idle_timeout: int = -1, # Negative values disable the idle timeout. + metadata: dict | None = None, + + ) -> None: + self._policy = policy + self._host = host + self._port = port + self._metadata = metadata or {} + self._idle_timeout = idle_timeout + self._last_active = time.time() + logging.getLogger("websockets.server").setLevel(logging.INFO) + + def serve_forever(self) -> None: + asyncio.run(self.run()) + + async def run(self): + async with websockets.asyncio.server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + ) as server: + if self._idle_timeout > 0: + await self._idle_watchdog(server) + else: + await server.serve_forever() + + async def _idle_watchdog(self, server): + """Close the server after the configured idle period.""" + while True: + await asyncio.sleep(5) + if time.time() - self._last_active > self._idle_timeout: + logging.info(f"Idle timeout ({self._idle_timeout}s) reached, shutting down server.") + server.close() + await server.wait_closed() + break + + async def _handler(self, websocket: websockets.asyncio.server.ServerConnection): + logging.info(f"Connection from {websocket.remote_address} opened") + packer = msgpack_numpy.Packer() + + await websocket.send(packer.pack(self._metadata)) + + while True: + try: + msg = msgpack_numpy.unpackb(await websocket.recv()) + self._last_active = time.time() + ret = self._route_message(msg) + await websocket.send(packer.pack(ret)) + except websockets.ConnectionClosed: + logging.info(f"Connection from {websocket.remote_address} closed") + break + except Exception: + await websocket.send(traceback.format_exc()) + await websocket.close( + code=websockets.frames.CloseCode.INTERNAL_ERROR, + reason="Internal server error. Traceback included in previous frame.", + ) + raise + + # route logic: recognize request from client + def _route_message(self, msg: dict) -> dict: + """ + Route rules (fault-tolerant): + - Supports messages of form: + {"type": "ping|init|infer|reset", "request_id": "...", "payload": {...}} + or a flat dict (will be treated as payload). + - Does NOT raise inside this function: all exceptions are caught and encoded in response. + """ + if not isinstance(msg, dict): + return { + "status": "error", + "ok": False, + "type": "inference_result", + "request_id": "default", + "error": {"message": "Payload must be a dict", "payload_type": str(type(msg))}, + } + + req_id = msg.get("request_id", "default") + mtype = msg.get("type", "infer") # default = infer + # The envelope fields are addressed to this router, not to the policy. + payload = {key: value for key, value in msg.items() if key not in ("type", "request_id")} + + # ping + if mtype == "ping": + return {"status": "ok", "ok": True, "type": "ping", "request_id": req_id} + + # infer --> framework.predict_action + elif mtype == "infer" or mtype == "predict_action": + try: + ouput_dict = self._policy.predict_action(**payload) + except Exception as e: + logging.exception("Policy inference error (request_id=%s)", req_id) + logging.exception(e) + + return { + "status": "error", + "ok": False, + "type": "inference_result", + "request_id": req_id, + "error": { + "message": str(e), + # "traceback": traceback.format_exc(), + }, + } + data = ouput_dict + return { + "status": "ok", + "ok": True, + "type": "inference_result", + "request_id": req_id, + "data": data, + } + + # unknow request type + else: + return { + "status": "error", + "ok": False, + "type": "unknown", + "request_id": req_id, + "error": {"message": f"Unsupported message type '{mtype}'"}, + } + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + # Example usage: + # policy = YourPolicyClass() # Replace with your actual policy class + # server = WebsocketPolicyServer(policy, host="localhost", port=10091) + # server.serve_forever() + raise NotImplementedError("This module is not intended to be run directly.") +# +# Instead, it should be imported and used in a server context. diff --git a/real_robot/source/deployment/model_server/xtrainer2_command.py b/real_robot/source/deployment/model_server/xtrainer2_command.py new file mode 100644 index 0000000000000000000000000000000000000000..029b2a322193350161e02e75d6bbc3bc5a9e7015 --- /dev/null +++ b/real_robot/source/deployment/model_server/xtrainer2_command.py @@ -0,0 +1,93 @@ +"""xtrainer2 command layout used between StarVLA and the real robot. + +The converted LeRobot dataset stores each arm as ``xyz + quaternion + gripper``. +StarVLA / PyTorch3D read that quaternion as ``w, x, y, z``. The robot HDF5 +streams and the robot controller use ``x, y, z, w``. + +The training transform already maps quaternion <-> rotation-6D. After +un-normalization the server therefore has 16-D **wxyz** commands. Sending +those 16 numbers to the robot without this reorder will execute the wrong +orientation. + +Do not confuse these widths: + +* 14 – joint angles in the source HDF5 (not used for this VLA) +* 16 – stored / robot EEF command (after this conversion, xyzw) +* 20 – model tensor after rotation-6D (horizon is a separate axis) +""" + +from __future__ import annotations + +import numpy as np + +from starVLA.dataloader.relative_pose import compose_quaternion + +QUATERNION_STARTS = (3, 11) +POSITION_SLICES = ((0, 3), (8, 11)) +ACTION_DIM = 16 +ROBOT_TYPE = "xtrainer2_bimanual_eef" +SOURCE_QUATERNION_FORMAT = "xyzw" +MODEL_QUATERNION_FORMAT = "wxyz" + +#: Value of ``DataConfig.robot_command_layout`` that selects the conversions in +#: this module. The policy server reads that attribute instead of matching on a +#: class name. +COMMAND_LAYOUT = "xtrainer2_bimanual_eef_xyzw" + + +def _as_float_array(values: np.ndarray) -> np.ndarray: + array = np.asarray(values, dtype=np.float32) + if array.shape[-1] != ACTION_DIM: + raise ValueError(f"expected last dimension {ACTION_DIM}, got {array.shape}") + return array + + +def wxyz_to_xyzw_block(block: np.ndarray) -> np.ndarray: + """``w, x, y, z`` -> ``x, y, z, w``.""" + return np.concatenate([block[..., 1:4], block[..., :1]], axis=-1) + + +def xyzw_to_wxyz_block(block: np.ndarray) -> np.ndarray: + """``x, y, z, w`` -> ``w, x, y, z``.""" + return np.concatenate([block[..., 3:4], block[..., :3]], axis=-1) + + +def model_actions_to_robot(values: np.ndarray) -> np.ndarray: + """Convert un-normalized StarVLA actions (wxyz) to robot commands (xyzw).""" + converted = np.array(_as_float_array(values), copy=True) + for start in QUATERNION_STARTS: + converted[..., start : start + 4] = wxyz_to_xyzw_block(converted[..., start : start + 4]) + return converted + + +def robot_state_to_model(values: np.ndarray) -> np.ndarray: + """Convert robot observations (xyzw) to the layout StarVLA was trained on.""" + converted = np.array(_as_float_array(values), copy=True) + for start in QUATERNION_STARTS: + converted[..., start : start + 4] = xyzw_to_wxyz_block(converted[..., start : start + 4]) + return converted + + +def compose_relative_eef_wxyz(relative_actions: np.ndarray, state_wxyz: np.ndarray) -> np.ndarray: + """Turn rel-mode 16-D wxyz actions into absolute 16-D wxyz poses. + + Positions add the current EE; rotations left-multiply the current + quaternion; grippers are already absolute and are left unchanged. + """ + relative = np.array(_as_float_array(relative_actions), copy=True) + state = np.asarray(state_wxyz, dtype=np.float32).reshape(-1, ACTION_DIM)[0] + absolute = relative + for start, end in POSITION_SLICES: + absolute[..., start:end] = relative[..., start:end] + state[start:end] + for start in QUATERNION_STARTS: + absolute[..., start : start + 4] = compose_quaternion( + state[start : start + 4], + relative[..., start : start + 4], + ) + return absolute + + +def is_relative_action_mode(action_mode: object | None) -> bool: + if action_mode is None: + return False + return str(action_mode).lower() in {"rel", "relative"} diff --git a/real_robot/source/examples/Xtrainer2/INFERENCE.md b/real_robot/source/examples/Xtrainer2/INFERENCE.md new file mode 100644 index 0000000000000000000000000000000000000000..7bcaa6089754e4e1c3df96c1229535492f692296 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/INFERENCE.md @@ -0,0 +1,4 @@ +# Xtrainer2 policy inference + +The server and client interface is described in [ROBOT.md](ROBOT.md). +Data layout and training configuration are described in [README.md](README.md). diff --git a/real_robot/source/examples/Xtrainer2/README.md b/real_robot/source/examples/Xtrainer2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..3d04e46b3b7858b47595cdd2a18df1952bb26aaf --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/README.md @@ -0,0 +1,26 @@ +# Xtrainer2 data adapter + +This directory contains optional StarVLA configurations and conversion helpers +for Xtrainer2 LeRobot data. It does not distribute robot data, checkpoints, or +on-robot control software. The ReTime command in the repository root is the +recommended entrypoint for proposal generation and training. + +The supplied configurations cover pack-phone with absolute end-effector actions +and pick-bottle with relative end-effector actions. For bimanual data, the source +action contains 16 values: position, quaternion, and gripper for each arm. +Rotation conversion expands the model action to 20 values. `pytorch3d` is +required for this conversion; install a build matching your Python, PyTorch, +and CUDA environment. The additional Python dependencies are listed in +`requirements-xtrainer2.txt`. + +Set dataset, model, statistics, training budget, and process count for your own +environment. Keep the trainer gradient accumulation equal to the Accelerate +configuration. Use a full checkpoint directory with `--start resume` when +optimizer and scheduler continuation is required. A model file with the default +`--start weights` initializes a new optimizer and scheduler. + +Policy serving uses the included `deployment/model_server` implementation. A +separate client must supply observations and consume actions using the same +schema as the selected data configuration. Client source and deployment +instructions are outside this artifact. See [ROBOT.md](ROBOT.md) for the data +and interface contract. diff --git a/real_robot/source/examples/Xtrainer2/ROBOT.md b/real_robot/source/examples/Xtrainer2/ROBOT.md new file mode 100644 index 0000000000000000000000000000000000000000..ff6472129420b78202bb2b70ff4700eb42d91b78 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/ROBOT.md @@ -0,0 +1,18 @@ +# Xtrainer2 policy interface + +The optional policy server is in `deployment/model_server`. Configure it with a +model checkpoint and the matching training configuration and dataset statistics. +The client is a separate component and is not distributed in this artifact. + +The client sends camera observations, proprioception, and a task instruction. +For the supplied bimanual configurations, its proprioception and actions use +16 absolute end-effector values: position, quaternion, and gripper for each arm. +The model internally uses 20 values after quaternion to rotation-6D conversion. +The client quaternion format is `xyzw`; dataset conversion and normalization +must follow the selected configuration. The pick-bottle configuration composes +relative actions on the server before returning absolute client actions. + +Check the model architecture, action dimension, normalization statistics, +camera names, and action chunk size against the checkpoint before deployment. +Only use a client whose message schema matches the server's documented policy +interface. Robot hardware and control software are outside this review package. diff --git a/real_robot/source/examples/Xtrainer2/train_files/build_pack_phone_support.py b/real_robot/source/examples/Xtrainer2/train_files/build_pack_phone_support.py new file mode 100644 index 0000000000000000000000000000000000000000..ba39a1579edf1c5a7963cb4c0868c73e2a9a3008 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/train_files/build_pack_phone_support.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Write frozen fingerprint and provenance for xtrainer2_pack_phone.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def _identity_files(dataset_root: Path) -> list[Path]: + files = list((dataset_root / "data").rglob("*.parquet")) + files.extend((dataset_root / "videos").rglob("*.mp4")) + meta_root = dataset_root / "meta" + if meta_root.exists(): + for pattern in ("*.json", "*.jsonl", "*.parquet"): + files.extend(meta_root.rglob(pattern)) + return sorted(set(files), key=lambda path: path.relative_to(dataset_root).as_posix()) + + +def compute_dataset_content_fingerprint(dataset_root: str | Path) -> dict[str, Any]: + root = Path(dataset_root) + entries = [] + total_bytes = 0 + for file_path in _identity_files(root): + relative_path = file_path.relative_to(root).as_posix() + file_size = file_path.stat().st_size + entries.append({"path": relative_path, "bytes": file_size}) + total_bytes += file_size + return { + "file_count": len(entries), + "total_bytes": total_bytes, + "files": entries, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--data-root", + type=Path, + default=Path("/path/to/data/lerobot"), + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + ) + args = parser.parse_args() + dataset_root = args.data_root / "xtrainer2_pack_phone" + output_dir = args.output_dir or (args.data_root / "xtrainer2_optr_support" / "pack_phone") + output_dir.mkdir(parents=True, exist_ok=True) + + info = json.loads((dataset_root / "meta" / "info.json").read_text(encoding="utf-8")) + report_path = dataset_root / "conversion_report.json" + report = json.loads(report_path.read_text(encoding="utf-8")) if report_path.is_file() else {} + + provenance = { + "dataset_name": "xtrainer2_pack_phone", + "codebase_version": info.get("codebase_version"), + "robot_type": info.get("robot_type"), + "fps": info.get("fps"), + "total_episodes": info.get("total_episodes"), + "total_frames": info.get("total_frames"), + "quaternion_format": info.get("quaternion_format"), + "source_quaternion_format": info.get("source_quaternion_format"), + "source_fingerprint": report.get("source_fingerprint"), + } + provenance_path = output_dir / "data_provenance.json" + provenance_path.write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + fingerprint = { + "datasets": { + "xtrainer2_pack_phone": compute_dataset_content_fingerprint(dataset_root), + } + } + fingerprint_path = output_dir / "content_fingerprints.json" + fingerprint_path.write_text(json.dumps(fingerprint) + "\n", encoding="utf-8") + + print(f"wrote {provenance_path}") + print(f"wrote {fingerprint_path}") + + +if __name__ == "__main__": + main() diff --git a/real_robot/source/examples/Xtrainer2/train_files/build_pick_bottle_support.py b/real_robot/source/examples/Xtrainer2/train_files/build_pick_bottle_support.py new file mode 100644 index 0000000000000000000000000000000000000000..987ce07fcb475c002c8a8f88349c182ba8db8ffb --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/train_files/build_pick_bottle_support.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Write frozen fingerprint and provenance for xtrainer2_pick_bottle.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def _identity_files(dataset_root: Path) -> list[Path]: + files = list((dataset_root / "data").rglob("*.parquet")) + files.extend((dataset_root / "videos").rglob("*.mp4")) + meta_root = dataset_root / "meta" + if meta_root.exists(): + for pattern in ("*.json", "*.jsonl", "*.parquet"): + files.extend(meta_root.rglob(pattern)) + return sorted(set(files), key=lambda path: path.relative_to(dataset_root).as_posix()) + + +def compute_dataset_content_fingerprint(dataset_root: str | Path) -> dict[str, Any]: + root = Path(dataset_root) + entries = [] + total_bytes = 0 + for file_path in _identity_files(root): + relative_path = file_path.relative_to(root).as_posix() + file_size = file_path.stat().st_size + entries.append({"path": relative_path, "bytes": file_size}) + total_bytes += file_size + return { + "file_count": len(entries), + "total_bytes": total_bytes, + "files": entries, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--data-root", + type=Path, + default=Path("/path/to/data/lerobot"), + ) + parser.add_argument( + "--output-dir", + type=Path, + default=None, + ) + args = parser.parse_args() + dataset_root = args.data_root / "xtrainer2_pick_bottle" + output_dir = args.output_dir or (args.data_root / "xtrainer2_optr_support" / "pick_bottle") + output_dir.mkdir(parents=True, exist_ok=True) + + info = json.loads((dataset_root / "meta" / "info.json").read_text(encoding="utf-8")) + report_path = dataset_root / "conversion_report.json" + report = json.loads(report_path.read_text(encoding="utf-8")) if report_path.is_file() else {} + + provenance = { + "dataset_name": "xtrainer2_pick_bottle", + "codebase_version": info.get("codebase_version"), + "robot_type": info.get("robot_type"), + "fps": info.get("fps"), + "total_episodes": info.get("total_episodes"), + "total_frames": info.get("total_frames"), + "total_tasks": info.get("total_tasks"), + "quaternion_format": info.get("quaternion_format"), + "source_quaternion_format": info.get("source_quaternion_format"), + "source_fingerprint": report.get("source_fingerprint"), + } + provenance_path = output_dir / "data_provenance.json" + provenance_path.write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + fingerprint = { + "datasets": { + "xtrainer2_pick_bottle": compute_dataset_content_fingerprint(dataset_root), + } + } + fingerprint_path = output_dir / "content_fingerprints.json" + fingerprint_path.write_text(json.dumps(fingerprint) + "\n", encoding="utf-8") + + print(f"wrote {provenance_path}") + print(f"wrote {fingerprint_path}") + + +if __name__ == "__main__": + main() diff --git a/real_robot/source/examples/Xtrainer2/train_files/data_registry/data_config.py b/real_robot/source/examples/Xtrainer2/train_files/data_registry/data_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8a40351a4dff552fb8851ab6ce92d3b991ff2ed6 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/train_files/data_registry/data_config.py @@ -0,0 +1,179 @@ +"""StarVLA data registry for xtrainer2 bimanual EEF datasets.""" + +from starVLA.dataloader.gr00t_lerobot.datasets import ModalityConfig +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag +from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform +from starVLA.dataloader.gr00t_lerobot.transform.state_action import ( + StateActionToTensor, + StateActionTransform, +) +from starVLA.dataloader.gr00t_lerobot.transform.video import ( + VideoColorJitter, + VideoCrop, + VideoResize, + VideoToNumpy, + VideoToTensor, +) + + +class Xtrainer2BimanualEefDataConfig: + """Three-camera bimanual EEF control at 10 Hz. + + LeRobot stores 16 raw dimensions per state/action in wxyz. Training + converts quaternions to rotation-6D (20 model dims) and inverts them + back to 16-D wxyz. With ``action_mode: rel`` the dataloader expresses + xyz and orientation relative to the current EE; the websocket server + composes those back to absolute poses and reorders wxyz -> xyzw. + See ``examples/Xtrainer2/README.md``. + """ + + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_position", + "state.left_rotation", + "state.left_gripper", + "state.right_position", + "state.right_rotation", + "state.right_gripper", + ] + action_keys = [ + "action.left_position", + "action.left_rotation", + "action.left_gripper", + "action.right_position", + "action.right_rotation", + "action.right_gripper", + ] + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + # Raw widths, matching ``meta/modality.json`` and the 16-D arrays in + # ``dataset_statistics.json``. 16 is not divisible by the 6 keys, so the + # policy server cannot infer the split and must be told. + state_key_dims = { + "state.left_position": 3, + "state.left_rotation": 4, + "state.left_gripper": 1, + "state.right_position": 3, + "state.right_rotation": 4, + "state.right_gripper": 1, + } + action_key_dims = { + "action.left_position": 3, + "action.left_rotation": 4, + "action.left_gripper": 1, + "action.right_position": 3, + "action.right_rotation": 4, + "action.right_gripper": 1, + } + # Source representation of every rotation key. The dataset carries this in + # ``meta/modality.json``; deployment rebuilds the transforms from the + # checkpoint alone, where that file is not available. + rotation_types = { + "state.left_rotation": "quaternion", + "state.right_rotation": "quaternion", + "action.left_rotation": "quaternion", + "action.right_rotation": "quaternion", + } + + # Deployment contract (read by deployment/model_server). The robot sends + # and receives 16 raw dims with xyzw quaternions; the model works in 20 + # normalized dims with wxyz. See deployment/model_server/xtrainer2_command.py. + robot_command_layout = "xtrainer2_bimanual_eef_xyzw" + server_normalizes_state = True + + def modality_config(self): + return { + "video": ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ), + "state": ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ), + "action": ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ), + "language": ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ), + } + + def transform(self): + rotation_keys = { + "state.left_rotation": "rotation_6d", + "state.right_rotation": "rotation_6d", + } + action_rotation_keys = { + "action.left_rotation": "rotation_6d", + "action.right_rotation": "rotation_6d", + } + state_normalization = { + key: "min_max" if key.endswith("_rotation") else "q99" + for key in self.state_keys + } + action_normalization = { + key: "min_max" if key.endswith("_rotation") else "q99" + for key in self.action_keys + } + return ComposedModalityTransform( + transforms=[ + VideoToTensor(apply_to=self.video_keys), + VideoCrop(apply_to=self.video_keys, scale=0.95), + VideoResize( + apply_to=self.video_keys, + height=224, + width=224, + interpolation="linear", + ), + VideoColorJitter( + apply_to=self.video_keys, + brightness=0.3, + contrast=0.4, + saturation=0.5, + hue=0.08, + ), + VideoToNumpy(apply_to=self.video_keys), + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes=state_normalization, + target_rotations=rotation_keys, + ), + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes=action_normalization, + target_rotations=action_rotation_keys, + ), + ] + ) + + +ROBOT_TYPE_CONFIG_MAP = { + "xtrainer2_bimanual_eef": Xtrainer2BimanualEefDataConfig(), +} + +ROBOT_TYPE_TO_EMBODIMENT_TAG = {} + +DATASET_NAMED_MIXTURES = { + "xtrainer2_pick_bottle": [ + ("xtrainer2_pick_bottle", 1.0, "xtrainer2_bimanual_eef"), + ], + "xtrainer2_pack_phone": [ + ("xtrainer2_pack_phone", 1.0, "xtrainer2_bimanual_eef"), + ], + "xtrainer2_pick_pack": [ + ("xtrainer2_pick_bottle", 1.0, "xtrainer2_bimanual_eef"), + ("xtrainer2_pack_phone", 1.0, "xtrainer2_bimanual_eef"), + ], +} diff --git a/real_robot/source/examples/Xtrainer2/train_files/modality.json b/real_robot/source/examples/Xtrainer2/train_files/modality.json new file mode 100644 index 0000000000000000000000000000000000000000..66fef208b71ddc0815a4203c51740e45e96d2b15 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/train_files/modality.json @@ -0,0 +1,26 @@ +{ + "state": { + "left_position": {"start": 0, "end": 3, "dtype": "float32", "absolute": true, "original_key": "observation.state"}, + "left_rotation": {"start": 3, "end": 7, "dtype": "float32", "absolute": true, "rotation_type": "quaternion", "original_key": "observation.state"}, + "left_gripper": {"start": 7, "end": 8, "dtype": "float32", "absolute": true, "original_key": "observation.state"}, + "right_position": {"start": 8, "end": 11, "dtype": "float32", "absolute": true, "original_key": "observation.state"}, + "right_rotation": {"start": 11, "end": 15, "dtype": "float32", "absolute": true, "rotation_type": "quaternion", "original_key": "observation.state"}, + "right_gripper": {"start": 15, "end": 16, "dtype": "float32", "absolute": true, "original_key": "observation.state"} + }, + "action": { + "left_position": {"start": 0, "end": 3, "dtype": "float32", "absolute": true, "original_key": "action"}, + "left_rotation": {"start": 3, "end": 7, "dtype": "float32", "absolute": true, "rotation_type": "quaternion", "original_key": "action"}, + "left_gripper": {"start": 7, "end": 8, "dtype": "float32", "absolute": true, "original_key": "action"}, + "right_position": {"start": 8, "end": 11, "dtype": "float32", "absolute": true, "original_key": "action"}, + "right_rotation": {"start": 11, "end": 15, "dtype": "float32", "absolute": true, "rotation_type": "quaternion", "original_key": "action"}, + "right_gripper": {"start": 15, "end": 16, "dtype": "float32", "absolute": true, "original_key": "action"} + }, + "video": { + "cam_high": {"original_key": "observation.images.cam_high"}, + "cam_left_wrist": {"original_key": "observation.images.cam_left_wrist"}, + "cam_right_wrist": {"original_key": "observation.images.cam_right_wrist"} + }, + "annotation": { + "human.action.task_description": {"original_key": "task_index"} + } +} diff --git a/real_robot/source/examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh b/real_robot/source/examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh new file mode 100644 index 0000000000000000000000000000000000000000..dc95f54cd96c56093e0ac3056178096238b9da07 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +# Start the StarVLA websocket policy server for xtrainer2. +# The robot PC talks to this process; do not clone RoboticsX into this repo. +# The on-robot client is a separate component. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" + +if [[ -f "${REPO_ROOT}/.env.xtrainer2-server" ]]; then + # Local GPU-server paths (gitignored). Sets OPTR_PYTHON / FLORENCE_MODEL_PATH. + # shellcheck disable=SC1091 + source "${REPO_ROOT}/.env.xtrainer2-server" +fi + +if [[ $# -lt 1 ]]; then + echo "Usage: bash examples/Xtrainer2/train_files/run_xtrainer2_policy_server.sh [gpu_id] [port]" >&2 + exit 1 +fi + +ckpt=$(readlink -f "$1") +gpu_id=${2:-0} +port=${3:-10093} +if [[ -z "${OPTR_PYTHON:-}" && -z "${STARVLA_PYTHON:-}" && -x "${REPO_ROOT}/.venv-xtrainer2/bin/python" ]]; then + python_bin="${REPO_ROOT}/.venv-xtrainer2/bin/python" +else + python_bin=${OPTR_PYTHON:-${STARVLA_PYTHON:-python}} +fi + +if [[ ! -f $ckpt ]]; then + echo "checkpoint not found: ${ckpt}" >&2 + exit 1 +fi + +# `read_mode_config` resolves config.yaml and dataset_statistics.json from the +# checkpoint's grandparent, i.e. it expects /checkpoints/.pt. A +# DeepSpeed run instead writes /checkpoints/steps_N_accelerate/pytorch_model.pt, +# which is one level deeper, so stage that layout out of the way of the run dir. +ckpt_dir=$(dirname "$ckpt") +if [[ -f "${ckpt_dir}/../config.yaml" ]]; then + run_dir=$(cd "${ckpt_dir}/.." && pwd) + serve_ckpt=$ckpt + stage_dir="" +elif [[ -f "${ckpt_dir}/../../config.yaml" ]]; then + run_dir=$(cd "${ckpt_dir}/../.." && pwd) + stage_dir="${TMPDIR:-/tmp}/starvla-serve/$(basename "$run_dir")/$(basename "$ckpt_dir")" + mkdir -p "${stage_dir}/checkpoints" + ln -sfn "${run_dir}/config.yaml" "${stage_dir}/config.yaml" + ln -sfn "${run_dir}/dataset_statistics.json" "${stage_dir}/dataset_statistics.json" + ln -sfn "$ckpt" "${stage_dir}/checkpoints/$(basename "$ckpt_dir")_pytorch_model.pt" + serve_ckpt="${stage_dir}/checkpoints/$(basename "$ckpt_dir")_pytorch_model.pt" + echo "[INFO] accelerate checkpoint layout; serving through ${stage_dir}" +else + echo "need config.yaml and dataset_statistics.json in the run dir above ${ckpt_dir}" >&2 + exit 1 +fi + +# Training config.yaml records the cluster Florence path. Rewrite it locally +# so from_pretrained does not look for /path/to/models/... . +if [[ -n "${FLORENCE_MODEL_PATH:-}" ]]; then + if [[ ! -d $FLORENCE_MODEL_PATH ]]; then + echo "FLORENCE_MODEL_PATH is not a directory: ${FLORENCE_MODEL_PATH}" >&2 + exit 1 + fi + if [[ -z $stage_dir ]]; then + stage_dir="${TMPDIR:-/tmp}/starvla-serve/$(basename "$run_dir")/local-florence" + mkdir -p "${stage_dir}/checkpoints" + ln -sfn "$ckpt" "${stage_dir}/checkpoints/$(basename "$ckpt")" + ln -sfn "${run_dir}/dataset_statistics.json" "${stage_dir}/dataset_statistics.json" + serve_ckpt="${stage_dir}/checkpoints/$(basename "$ckpt")" + fi + FLORENCE_MODEL_PATH="$FLORENCE_MODEL_PATH" "$python_bin" - "${run_dir}/config.yaml" "${stage_dir}/config.yaml" <<'PY' +import os +import sys + +from omegaconf import OmegaConf + +src, dst = sys.argv[1], sys.argv[2] +path = os.environ["FLORENCE_MODEL_PATH"] +if "florence" not in path.lower(): + raise SystemExit(f"FLORENCE_MODEL_PATH must contain 'florence': {path}") +cfg = OmegaConf.load(src) +cfg.framework.qwenvl.base_vlm = path +OmegaConf.save(cfg, dst) +print(f"[INFO] rewrote base_vlm -> {path}") +PY +fi + +if [[ ! -f $run_dir/dataset_statistics.json ]]; then + echo "missing ${run_dir}/dataset_statistics.json" >&2 + exit 1 +fi + +echo "[INFO] xtrainer2 policy server" +echo "[INFO] checkpoint: ${ckpt}" +echo "[INFO] run dir: ${run_dir}" +echo "[INFO] gpu: ${gpu_id} port: ${port}" +echo "[INFO] client contract: 3 images (cam_high, cam_left_wrist, cam_right_wrist)" +echo "[INFO] HWC RGB uint8, 16-D robot state and 16-D actions, quaternions xyzw" +echo "[INFO] instruction must be English: Pick up the phone, place it in the box, and close the lid." + +# Local Florence weights; do not send Hub traffic through a lab HTTP proxy. +unset HF_ENDPOINT || true +idle_timeout="${OPTR_IDLE_TIMEOUT:--1}" +CUDA_VISIBLE_DEVICES=$gpu_id "$python_bin" \ + "$REPO_ROOT/deployment/model_server/server_policy.py" \ + --ckpt_path "$serve_ckpt" \ + --port "$port" \ + --use_bf16 \ + --idle_timeout "$idle_timeout" diff --git a/real_robot/source/examples/Xtrainer2/train_files/run_xtrainer2_train.sh b/real_robot/source/examples/Xtrainer2/train_files/run_xtrainer2_train.sh new file mode 100644 index 0000000000000000000000000000000000000000..9664c487033c5991098b6e1813b0ab4bf54d36b8 --- /dev/null +++ b/real_robot/source/examples/Xtrainer2/train_files/run_xtrainer2_train.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 2 ]]; then + echo "usage: $0 " >&2 + exit 64 +fi + +project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)}" +python="${OPTR_PYTHON:-python}" +accelerate_config="${OPTR_ACCELERATE_CONFIG:-${project}/configs/deepspeed_florence_8gpu.yaml}" +num_processes="${OPTR_NUM_PROCESSES:-8}" +config=$1 +run_id=$2 + +if [[ ! -f "$config" && -f "${project}/${config}" ]]; then + config="${project}/${config}" +fi +if [[ ! -f "$config" ]]; then + echo "error: training config not found: ${config}" >&2 + exit 66 +fi + +# The ${oc.env:...} refs in the training yaml have no defaults; OmegaConf +# resolves them eagerly and would abort on all ranks. Collect names from the +# file so pack_phone and pick_bottle (different provenance vars) share this +# launcher. +mapfile -t required_vars < <("$python" - "$config" <<'PY' +import re +import sys +from pathlib import Path + +text = Path(sys.argv[1]).read_text(encoding="utf-8") +names = sorted(set(re.findall(r"\$\{oc\.env:([A-Za-z_][A-Za-z0-9_]*)\}", text))) +if not names: + sys.exit("error: no ${oc.env:...} interpolations found in " + sys.argv[1]) +print("\n".join(names)) +PY +) + +missing=() +for var in "${required_vars[@]}"; do + if [[ -z "${!var:-}" ]]; then + missing+=("$var") + fi +done +if (( ${#missing[@]} )); then + echo "error: required environment variables are unset: ${missing[*]}" >&2 + echo "see examples/Xtrainer2/README.md for the export block" >&2 + exit 78 +fi + +if [[ -n "${FLORENCE_MODEL_PATH:-}" ]]; then + if [[ "${FLORENCE_MODEL_PATH,,}" != *florence* ]]; then + echo "error: FLORENCE_MODEL_PATH must contain 'florence' or the VLM dispatch" >&2 + echo " raises NotImplementedError: ${FLORENCE_MODEL_PATH}" >&2 + exit 78 + fi + if [[ ! -d "$FLORENCE_MODEL_PATH" ]]; then + echo "error: FLORENCE_MODEL_PATH is not a directory: ${FLORENCE_MODEL_PATH}" >&2 + exit 66 + fi +fi + +for var in "${required_vars[@]}"; do + value="${!var}" + case "$var" in + FLORENCE_MODEL_PATH|OPTR_RUN_ROOT|XTRAINER2_DATA_ROOT) continue ;; + esac + if [[ "$var" == *_PATH || "$var" == *_DIR ]]; then + if [[ ! -e "$value" ]]; then + echo "error: support file missing: ${value} (${var})" >&2 + echo " run the matching examples/Xtrainer2/train_files/build_*_support.py first" >&2 + exit 66 + fi + fi +done + +if ! command -v "$python" >/dev/null 2>&1 && [[ ! -x "$python" ]]; then + echo "error: python interpreter not found: ${python}" >&2 + echo " set OPTR_PYTHON to an interpreter with deepspeed and pytorch3d" >&2 + exit 69 +fi + +if [[ ! -f "$accelerate_config" ]]; then + echo "error: accelerate config not found: ${accelerate_config}" >&2 + exit 66 +fi + +cd "$project" +export PYTHONPATH="${project}:${PYTHONPATH:-}" +export WANDB_MODE="${WANDB_MODE:-offline}" +"$python" -m accelerate.commands.launch \ + --config_file "$accelerate_config" \ + --num_processes "$num_processes" \ + starVLA/training/train_starvla.py \ + --config_yaml "$config" \ + --run_id "$run_id" diff --git a/real_robot/source/licenses/Diffusion-Policy-LICENSE b/real_robot/source/licenses/Diffusion-Policy-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a2eb6cb62c1472fa2de9b6f1f6e0f97af8ba96e7 --- /dev/null +++ b/real_robot/source/licenses/Diffusion-Policy-LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Columbia Artificial Intelligence and Robotics Lab + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/real_robot/source/licenses/Isaac-GR00T-LICENSE b/real_robot/source/licenses/Isaac-GR00T-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..c6b51dc2713a75b943191f5767865b372a3cc80e --- /dev/null +++ b/real_robot/source/licenses/Isaac-GR00T-LICENSE @@ -0,0 +1,190 @@ + 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 the 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 the 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 any 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 + + Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. 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. diff --git a/real_robot/source/licenses/MODIFIED_FILES.txt b/real_robot/source/licenses/MODIFIED_FILES.txt new file mode 100644 index 0000000000000000000000000000000000000000..ecc60c62dcd8a5662845e1fc705ca32821b93fa9 --- /dev/null +++ b/real_robot/source/licenses/MODIFIED_FILES.txt @@ -0,0 +1,17 @@ +starVLA/dataloader/lerobot_datasets.py +starVLA/dataloader/__init__.py +starVLA/training/train_starvla.py +starVLA/training/trainer_utils/overwatch.py +starVLA/training/trainer_utils/trainer_tools.py +starVLA/model/framework/base_framework.py +starVLA/model/framework/share_tools.py +starVLA/model/modules/action_model/GR00T_ActionHeader.py +starVLA/model/modules/action_model/__init__.py +starVLA/model/modules/vlm/Florence2.py +starVLA/model/modules/vlm/QWen3.py +starVLA/model/modules/vlm/__init__.py +starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py +starVLA/model/framework/VLM4A/QwenGR00T.py +starVLA/dataloader/gr00t_lerobot/datasets.py +starVLA/dataloader/gr00t_lerobot/data_config.py +starVLA/dataloader/gr00t_lerobot/transform/state_action.py diff --git a/real_robot/source/licenses/OpenVLA-LICENSE b/real_robot/source/licenses/OpenVLA-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..04f26a7d0bfdeb89ce54681c695a40df21aa6d96 --- /dev/null +++ b/real_robot/source/licenses/OpenVLA-LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/real_robot/source/licenses/X-VLA-LICENSE b/real_robot/source/licenses/X-VLA-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/real_robot/source/licenses/X-VLA-LICENSE @@ -0,0 +1,201 @@ + 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 [yyyy] [name of copyright owner] + + 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/real_robot/source/licenses/msgpack-numpy-LICENSE b/real_robot/source/licenses/msgpack-numpy-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..e74fcfc1184b4e8c57827b715554111f5f197f8c --- /dev/null +++ b/real_robot/source/licenses/msgpack-numpy-LICENSE @@ -0,0 +1,35 @@ + +License +======= + +Copyright (c) 2013-2022, Lev E. Givon. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are +met: + +* Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. +* Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. +* Neither the name of Lev E. Givon nor the names of any + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/real_robot/source/pyproject.toml b/real_robot/source/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..7d9e3b5dd2c8082b6af9061a244aa5667d425134 --- /dev/null +++ b/real_robot/source/pyproject.toml @@ -0,0 +1,81 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + +[project] +name = "retime-starvla" +description = "ReTime temporal sampling plugin and compatible StarVLA backend." +version = "0.1.0" +readme = "README.md" +requires-python = ">=3.10" +keywords = ["vision-language-action models", "robot learning", "importance sampling"] +license = {file = "LICENSE"} +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3 :: Only", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = ["omegaconf>=2.3,<3", "numpy", "scipy", "torch"] + +[project.scripts] +retime = "retime.cli:main" +retime-pretrain = "retime.pretrain:main" + +[project.optional-dependencies] +dev = [ + "black>=24.2.0", + "ruff>=0.2.2", +] + + +[tool.setuptools] +license-files = ["LICENSE", "NOTICE", "licenses/*", "third_party/frameskip/LICENSE"] + +[tool.setuptools.packages.find] +where = ["."] +exclude = [ + "tests", + "tests.*", + "cache", + "cache.*", + "playground", + "playground.*", + "results", + "results.*", + "checkpoints", + "checkpoints.*", + "scripts", + "scripts.*", + "eval", + "eval.*", + "assets", + "assets.*", + ".vscode", + "*.egg-info", + "**__pycache__**", + "*.pyc", + "*.pyo", + "starVLA.egg-info", + "starVLA.egg-info.*", + "starVLA.config", + "starVLA.config.*", +] + +[tool.black] +line-length = 121 +target-version = ["py310"] +preview = true + +[tool.ruff] +line-length = 121 +target-version = "py310" + +[tool.ruff.lint] +select = ["A", "B", "E", "F", "I", "RUF", "W"] +ignore = ["F722"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["E402", "F401"] diff --git a/real_robot/source/requirements-qwen.txt b/real_robot/source/requirements-qwen.txt new file mode 100644 index 0000000000000000000000000000000000000000..a47928ddff6e57285fea3670578186fa4fc83d05 --- /dev/null +++ b/real_robot/source/requirements-qwen.txt @@ -0,0 +1,9 @@ +# Optional dependencies for the supplied Qwen-GR00T configurations. +# Install requirements.txt first. Qwen Bridge and Qwen RoboTwin are included in +# the GPU integration test documented in VALIDATION.md. + +-r requirements.txt +bitsandbytes==0.50.0 +packaging +psutil +ninja diff --git a/real_robot/source/requirements-xtrainer2.txt b/real_robot/source/requirements-xtrainer2.txt new file mode 100644 index 0000000000000000000000000000000000000000..42bc94da3e9ec1ee73b217998addf674b431b0c4 --- /dev/null +++ b/real_robot/source/requirements-xtrainer2.txt @@ -0,0 +1,10 @@ +# Optional dependencies for Xtrainer2 data processing and policy serving. +# Install requirements.txt first. This extension was not included in the +# Florence Bridge GPU integration test documented in VALIDATION.md. +# +# Xtrainer2 action_dim=20 requires a PyTorch3D build compatible with the +# installed Python, PyTorch, and CUDA versions. Install it separately from +# its official distribution for your platform. + +-r requirements.txt +pytest diff --git a/real_robot/source/requirements.txt b/real_robot/source/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..a244d2011e3410e557a74cacc6a753a4c4161237 --- /dev/null +++ b/real_robot/source/requirements.txt @@ -0,0 +1,34 @@ +transformers==4.57.0 +accelerate==1.13.0 +tiktoken +einops +transformers_stream_generator==0.0.4 +scipy +torch==2.7.1 +torchvision==0.22.1 +setuptools==80.9.0 +pillow +tensorboard +matplotlib +albumentations==1.4.18 +decord==0.6.0 +pydantic==2.10.6 +pyarrow==14.0.1 +fastparquet==2024.11.0 +av==12.3.0 +numpydantic==1.6.9 +deepspeed==0.19.2 +qwen-vl-utils +omegaconf +numpy==1.26.4 +wandb +rich +diffusers==0.33.1 +timm +tyro +websockets==15.0.1 +msgpack==1.1.0 +tdigest==0.5.2.2 + +# Qwen configurations also require requirements-qwen.txt and a compatible +# Flash Attention 2 installation; see README.md. diff --git a/real_robot/source/retime/__init__.py b/real_robot/source/retime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..392a119edfe7a378da8aaa40370f8e9376327511 --- /dev/null +++ b/real_robot/source/retime/__init__.py @@ -0,0 +1,6 @@ +"""ReTime: objective-preserving temporal sampling for StarVLA. + +Importing the plugin or its CLI does not initialize a model or CUDA. +""" + +__version__ = "0.1.0" diff --git a/real_robot/source/retime/__main__.py b/real_robot/source/retime/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..bfdcd0c1158dff4fcc25157b317419f0a89dd33f --- /dev/null +++ b/real_robot/source/retime/__main__.py @@ -0,0 +1,4 @@ +from .cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/real_robot/source/retime/checkpoint.py b/real_robot/source/retime/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..7dc9a9364d51c0c8bb7ed5753eb7c729099f62da --- /dev/null +++ b/real_robot/source/retime/checkpoint.py @@ -0,0 +1,19 @@ +"""Load full policy weights before or after DeepSpeed partitions parameters.""" + + +def load_policy_state_dict(model, state_dict, *, strict=True): + """Every distributed rank must call this with the same state dictionary. + + ZeRO-3 initialization replaces parameter storage with empty shards. Gather + those parameters while copying weights, then broadcast and repartition on + context exit. Loading on every rank also initializes replicated buffers. + Ordinary models do not need DeepSpeed installed. + """ + partitioned = [parameter for parameter in model.parameters() if hasattr(parameter, "ds_id")] + if not partitioned: + return model.load_state_dict(state_dict, strict=strict) + + from deepspeed import zero + + with zero.GatheredParameters(partitioned, modifier_rank=0): + return model.load_state_dict(state_dict, strict=strict) diff --git a/real_robot/source/retime/cli.py b/real_robot/source/retime/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..dc17c0f353c0c1ea1e76f0f5b1afdfbc92ba6f66 --- /dev/null +++ b/real_robot/source/retime/cli.py @@ -0,0 +1,272 @@ +"""One-command adapter to the accompanying StarVLA training backend. + +The CLI prepares temporal proposals and dispatches StarVLA training. +The sampler and ordinary importance correction are unchanged. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from pathlib import Path +import subprocess +import sys + + +def parser() -> argparse.ArgumentParser: + result = argparse.ArgumentParser(description=__doc__) + result.add_argument("--config", type=Path, required=True) + result.add_argument("--base-policy", type=Path, required=True, + help="Complete Accelerate checkpoint directory, committed single-process model file, or model file with --start weights.") + result.add_argument("--start", choices=("resume", "weights"), default="weights", + help="resume restores the full state; weights explicitly starts a new optimizer/schedule.") + result.add_argument("--alpha", type=float, default=0.2) + result.add_argument("--K", "--k", dest="k", type=int, default=8) + result.add_argument("--mode", nargs="+", choices=("retime", "uniform", "shuffle"), + default=["retime"]) + result.add_argument("--manifest", type=Path) + result.add_argument("--data-root", type=Path, help="Override the dataset root in the config") + result.add_argument("--data-mix", type=str, help="Override the named dataset mixture in the config") + result.add_argument("--vlm-path", type=Path, help="Override the Florence/Qwen backbone location") + result.add_argument("--statistics-path", type=Path, help="Override precomputed dataset statistics, if required") + result.add_argument( + "--omit-data-identity-files", + action="store_true", + help="Use the manifest's computed dataset identity without separate provenance/fingerprint files.", + ) + result.add_argument("--epsilon", type=float, default=0.1) + result.add_argument("--shuffle-seed", type=int, default=12345) + result.add_argument("--accelerate-config", type=Path, required=True) + result.add_argument("--output", type=Path, required=True) + result.add_argument("--enable-wandb", action="store_true") + result.add_argument("--backend-root", type=Path, default=Path(__file__).resolve().parents[1], + help="ReTime-enabled StarVLA source tree shipped with this artifact.") + result.add_argument("--dry-run", action="store_true", + help="Validate local inputs and print commands; do not score, train, or write files.") + return result + + +def _file(path: Path, label: str) -> Path: + path = path.expanduser().resolve() + if not path.is_file(): + raise ValueError(f"{label} is not a file: {path}") + return path + + +def _directory(path: Path, label: str) -> Path: + path = path.expanduser().resolve() + if not path.is_dir(): + raise ValueError(f"{label} is not a directory: {path}") + return path + + +def _native_commit_path(model_path: Path) -> Path: + for suffix in ("_pytorch_model.pt", "_model.safetensors"): + if model_path.name.endswith(suffix): + return model_path.with_name( + model_path.name[: -len(suffix)] + "_checkpoint.json" + ) + raise ValueError( + "A single-process resume checkpoint must end in " + "_pytorch_model.pt or _model.safetensors" + ) + + +def _validate_resume_checkpoint(base: Path, world_size: int) -> tuple[Path, dict]: + if base.is_dir(): + model = _file(base / "pytorch_model.pt", "base-policy model") + commit = json.loads( + _file(base / "checkpoint.json", "checkpoint commit").read_text( + encoding="utf-8" + ) + ) + if int(commit.get("world_size", -1)) != world_size: + raise ValueError("Base checkpoint world size differs from the launch contract") + entries = commit.get("files", []) + if not entries or any( + not (base / entry["name"]).is_file() + or (base / entry["name"]).stat().st_size != int(entry["bytes"]) + for entry in entries + ): + raise ValueError("Base Accelerate checkpoint is incomplete") + return model, commit + + model = _file(base, "base-policy model") + if world_size != 1: + raise ValueError( + "A committed single-process checkpoint requires a one-process launch" + ) + commit_path = _native_commit_path(model) + commit = json.loads(_file(commit_path, "checkpoint commit").read_text(encoding="utf-8")) + files = commit.get("files", {}) + for label in ("model", "training_state"): + entry = files.get(label, {}) + path = model.parent / str(entry.get("name", "")) + if not path.is_file() or path.stat().st_size != int(entry.get("bytes", -1)): + raise ValueError(f"Base single-process checkpoint has an invalid {label} file") + if files["model"]["name"] != model.name: + raise ValueError("Checkpoint commit does not describe --base-policy") + return model, commit + + +def build_plan(args: argparse.Namespace) -> dict: + from omegaconf import OmegaConf + + if not math.isfinite(args.alpha) or not 0 <= args.alpha < 1: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if args.k < 1 or not math.isfinite(args.epsilon) or args.epsilon <= 0: + raise ValueError("K and epsilon must be positive") + if len(set(args.mode)) != len(args.mode): + raise ValueError("Each mode can be requested only once") + root = args.backend_root.expanduser().resolve() + config = _file(args.config, "config") + launch_config = _file(args.accelerate_config, "accelerate config") + cfg = OmegaConf.load(config) + if not args.enable_wandb: + cfg.wandb_project = None + launch = OmegaConf.load(launch_config) + output = args.output.expanduser().resolve() + cfg.run_root_dir = str(output / "runs") + cfg.output_dir = str(output / "prepare") + if args.data_root: + cfg.datasets.vla_data.data_root_dir = str(_directory(args.data_root, "dataset root")) + if args.data_mix: + cfg.datasets.vla_data.data_mix = args.data_mix + if args.vlm_path: + cfg.framework.qwenvl.base_vlm = str(_directory(args.vlm_path, "VLM backbone")) + if args.statistics_path: + cfg.datasets.vla_data.precomputed_statistics_path = str(_file(args.statistics_path, "statistics")) + if args.omit_data_identity_files: + cfg.datasets.vla_data.dataset_provenance_path = None + cfg.datasets.vla_data.dataset_content_fingerprint_path = None + # Model-only initialization derives support from the current dataset. + if args.start == "weights": + cfg.trainer.is_resume = False + cfg.trainer.resume_accelerate_checkpoint = None + cfg.trainer.resume_checkpoint_integrity_manifest = None + if OmegaConf.select(cfg, "trainer.initial_completed_steps") is not None: + cfg.trainer.initial_completed_steps = 0 + if args.manifest is None: + cfg.datasets.vla_data.valid_anchor_manifest_path = None + cfg.datasets.vla_data.dataset_provenance_path = None + cfg.datasets.vla_data.dataset_content_fingerprint_path = None + cfg.datasets.vla_data.temporal_proposal_path = None + if args.manifest is not None: + cfg.datasets.vla_data.valid_anchor_manifest_path = str(args.manifest.expanduser().resolve()) + if OmegaConf.select(cfg, "framework.name") != "QwenGR00T": + raise ValueError("This adapter supports the supplied Florence/Qwen GR00T configs only") + if not OmegaConf.select(cfg, "datasets.vla_data.traceable_sampling", default=False): + raise ValueError("ReTime requires traceable_sampling and a frozen valid-anchor manifest") + world_size = int(launch.num_processes) + if int(OmegaConf.select(launch, "num_machines", default=1)) != 1: + raise ValueError("The one-command launcher currently supports one machine") + if world_size != int(cfg.datasets.vla_data.expected_world_size): + raise ValueError("Accelerate process count differs from the dataset training contract") + accumulation = int(cfg.trainer.gradient_accumulation_steps) + ds_accumulation = OmegaConf.select(launch, "deepspeed_config.gradient_accumulation_steps") + if ds_accumulation is not None and int(ds_accumulation) != accumulation: + raise ValueError("Accelerate and trainer gradient accumulation differ") + batch = int(cfg.datasets.vla_data.per_device_batch_size) + if min(world_size, accumulation, batch) < 1: + raise ValueError("World size, accumulation and batch size must be positive") + configured_manifest = OmegaConf.select(cfg, "datasets.vla_data.valid_anchor_manifest_path") + manifest = args.manifest or (Path(configured_manifest) if configured_manifest else None) + prepare_manifest = manifest is None + if prepare_manifest and args.start == "resume": + raise ValueError("Exact resume requires the parent run's manifest; use --manifest") + manifest = output / "proposal" / "manifest.json" if prepare_manifest else _file(manifest, "manifest") + if configured_manifest and Path(configured_manifest).expanduser().resolve() != manifest: + raise ValueError("--manifest differs from the config manifest") + base = args.base_policy.expanduser().resolve() + if args.start == "resume": + checkpoint, _ = _validate_resume_checkpoint(base, world_size) + else: + checkpoint = _file(base / "pytorch_model.pt" if base.is_dir() else base, "base-policy model") + cfg.trainer.pretrained_checkpoint = str(checkpoint) if args.start == "weights" else None + cfg.trainer.resume_accelerate_checkpoint = str(base) if args.start == "resume" else None + cfg.trainer.is_resume = args.start == "resume" + modes = list(args.mode) + for mode in modes: + if (output / "runs" / mode).exists(): + raise ValueError(f"Training output already exists: {output / 'runs' / mode}; choose a fresh --output") + commands = [] + cfg.datasets.vla_data.valid_anchor_manifest_path = str(manifest) + # Materialize all paths once so every arm and subprocess uses the same config. + config_text = OmegaConf.to_yaml(cfg, resolve=True) + config = output / "config.yaml" + if prepare_manifest: + commands.append([sys.executable, "-m", "retime.prepare", "--config", str(config), "--output", str(manifest)]) + + def script(name: str, *arguments: str) -> None: + module = { + "build_residual_temporal_proposal.py": "retime.commands.proposal", + "build_shuffle_temporal_proposal.py": "retime.commands.shuffle", + }[name] + commands.append([sys.executable, "-m", module, *map(str, arguments)]) + + residual = output / "proposal" / "retime.json" + shuffled = output / "proposal" / "shuffle.json" + if any(mode != "uniform" for mode in modes): + script("build_residual_temporal_proposal.py", "--config", config, "--checkpoint", checkpoint, + "--manifest", manifest, "--cache-dir", output / "cache", + "--output", residual, "--grid-size", args.k, "--alpha", args.alpha, "--epsilon", args.epsilon) + if "shuffle" in modes: + script("build_shuffle_temporal_proposal.py", "--manifest", manifest, + "--residual-proposal", residual, "--permutation-seed", args.shuffle_seed, + "--output", shuffled) + training = _file(root / "starVLA" / "training" / "train_starvla.py", "training backend") + for mode in modes: + proposal = {"uniform": None, "retime": str(residual), "shuffle": str(shuffled)}[mode] + overrides = { + "run_id": mode, + "run_root_dir": str(output / "runs"), + "output_dir": str(output / "runs" / mode), + "datasets.vla_data.valid_anchor_manifest_path": str(manifest), + "datasets.vla_data.temporal_proposal_path": proposal, + "trainer.is_resume": args.start == "resume", + "trainer.allow_model_only_resume": False, + "trainer.resume_accelerate_checkpoint": str(base) if args.start == "resume" else None, + "trainer.pretrained_checkpoint": str(checkpoint) if args.start == "weights" else None, + } + commands.append([ + sys.executable, "-m", "accelerate.commands.launch", "--config_file", str(launch_config), + "--num_processes", str(world_size), str(training), "--config_yaml", str(config), + *[f"--{key}={json.dumps(value)}" for key, value in overrides.items()], + ]) + return { + "backend_root": str(root), "output": str(output), "base_policy": str(base), + "start": args.start, "alpha": args.alpha, "K": args.k, "modes": modes, + "global_batch_size": batch * world_size * accumulation, + "config_text": config_text, + "commands": commands, + } + + +def main(argv: list[str] | None = None) -> int: + arguments = parser() + args = arguments.parse_args(argv) + try: + plan = build_plan(args) + except (ValueError, OSError, ImportError, KeyError) as error: + arguments.error(str(error)) + print(json.dumps({key: value for key, value in plan.items() if key != "config_text"}, indent=2), flush=True) + if args.dry_run: + return 0 + output = Path(plan["output"]) + output.mkdir(parents=True, exist_ok=True) + (output / "proposal").mkdir(exist_ok=True) + (output / "config.yaml").write_text(plan["config_text"], encoding="utf-8") + (output / "plan.json").write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8") + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join(filter(None, [plan["backend_root"], env.get("PYTHONPATH")])) + env.setdefault("WANDB_MODE", "offline") + env["PYTHONUNBUFFERED"] = "1" + for command in plan["commands"]: + subprocess.run(command, cwd=plan["backend_root"], env=env, check=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/real_robot/source/retime/commands/__init__.py b/real_robot/source/retime/commands/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..35c2c9c504f3a0b98708a2fb9fe2a14b9f0340c6 --- /dev/null +++ b/real_robot/source/retime/commands/__init__.py @@ -0,0 +1 @@ +"""ReTime proposal-generation commands.""" diff --git a/real_robot/source/retime/commands/proposal.py b/real_robot/source/retime/commands/proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..692c399df24fddf51d004a60bbd5af14a686d4e4 --- /dev/null +++ b/real_robot/source/retime/commands/proposal.py @@ -0,0 +1,695 @@ +"""Build a resumable, grid-scored residual temporal proposal. + +The expensive frozen-policy residual is evaluated at a fixed quantile grid +inside every episode. Scores for the remaining anchors are filled by +deterministic linear interpolation, ranked within that episode, and converted +to epsilon-floored masses. ``TemporalProposal`` then applies the safe mixture +and exact ``p/q`` correction used by training. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import random +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from omegaconf import OmegaConf + +from retime.proposal import TemporalProposal +from retime.scoring.residual_cache import ( + interpolate_grid_scores, + percentile_intensity_masses, + quantile_anchor_positions, +) +from retime.scoring.starvla_likelihood import ( + score_starvla_flow_matching_residual, +) + + +SCHEMA_VERSION = 1 + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + with temporary.open("w", encoding="utf-8") as output: + json.dump(payload, output, indent=2, sort_keys=True) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _set_random_seed(seed: int) -> None: + bounded_seed = int(seed) % (2**32) + random.seed(bounded_seed) + np.random.seed(bounded_seed) + torch.manual_seed(bounded_seed) + torch.cuda.manual_seed_all(bounded_seed) + + +def _json_scalar(value: Any) -> Any: + return value.item() if isinstance(value, np.generic) else value + + +def _shard_path(cache_dir: Path, dataset_index: int, episode_index: int) -> Path: + return cache_dir / "shards" / ( + f"dataset_{dataset_index:02d}_episode_{episode_index:06d}.json" + ) + + +def _planned_grid_count(manifest, grid_size: int) -> int: + return sum( + min(len(episode.anchors), grid_size) + for dataset in manifest.datasets + for episode in dataset.episodes + ) + + +def _episode_belongs_to_shard( + global_episode_index: int, + *, + shard_index: int, + num_shards: int, +) -> bool: + if num_shards <= 0: + raise ValueError("num_shards must be positive") + if not 0 <= shard_index < num_shards: + raise ValueError("shard_index must satisfy 0 <= shard_index < num_shards") + return int(global_episode_index) % int(num_shards) == int(shard_index) + + +def _worker_progress_path(cache_dir: Path, shard_index: int, num_shards: int) -> Path: + return cache_dir / "workers" / f"shard_{shard_index:02d}_of_{num_shards:02d}.json" + + +def _load_valid_shard( + path: Path, + *, + identity: dict[str, Any], + dataset_index: int, + episode_index: int, + trajectory_id: Any, + positions: np.ndarray, + base_indices: list[int], + estimator_count: int, +) -> dict[str, Any]: + payload = json.loads(path.read_text(encoding="utf-8")) + expected = { + "identity": identity, + "dataset_index": dataset_index, + "episode_index": episode_index, + "trajectory_id": _json_scalar(trajectory_id), + "anchor_positions": positions.tolist(), + "base_indices": base_indices, + } + for field, value in expected.items(): + if payload.get(field) != value: + raise ValueError(f"Residual cache shard mismatch for {field}: {path}") + estimators = np.asarray(payload.get("estimator_residuals"), dtype=np.float64) + if estimators.shape != (len(positions), estimator_count): + raise ValueError(f"Residual cache shard has wrong estimator shape: {path}") + if not np.isfinite(estimators).all(): + raise ValueError(f"Residual cache shard contains non-finite values: {path}") + scores = np.asarray(payload.get("grid_scores"), dtype=np.float64) + if scores.shape != (len(positions),) or not np.isfinite(scores).all(): + raise ValueError(f"Residual cache shard contains invalid grid scores: {path}") + return payload + + +def _load_examples( + source_dataset, + episode, + positions: np.ndarray, + *, + dataset_index: int, + episode_index: int, + transform_seed_base: int, +) -> tuple[list[dict], list[int], list[int]]: + examples = [] + base_indices = [] + transform_seeds = [] + for anchor_position in positions: + base_index = int(episode.anchors[int(anchor_position)]) + transform_seed = ( + transform_seed_base + + dataset_index * 10_000_019 + + episode_index * 10_007 + + int(anchor_position) + ) + _set_random_seed(transform_seed) + raw_data = source_dataset.get_step_data( + episode.trajectory_id, + base_index, + ) + transformed = source_dataset.transforms(raw_data) + examples.append(source_dataset._pack_sample(transformed)) + base_indices.append(base_index) + transform_seeds.append(transform_seed) + return examples, base_indices, transform_seeds + + +def _score_examples( + model, + examples: list[dict], + *, + estimator_seeds: list[int], + samples_per_estimator: int, + batch_size: int, +) -> np.ndarray: + chunks = [] + with torch.inference_mode(): + for start in range(0, len(examples), batch_size): + batch = examples[start : start + batch_size] + result = score_starvla_flow_matching_residual( + model, + batch, + estimator_seeds=estimator_seeds, + samples_per_estimator=samples_per_estimator, + ) + chunks.append(result.estimator_residuals.transpose(0, 1)) + return torch.cat(chunks, dim=0).float().cpu().numpy() + + +def _build_identity(args, manifest) -> dict[str, Any]: + checkpoint = args.checkpoint.resolve() + digest = hashlib.sha256() + with checkpoint.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return { + "schema_version": SCHEMA_VERSION, + "checkpoint": {"sha256": digest.hexdigest(), "bytes": checkpoint.stat().st_size}, + "config_sha256": hashlib.sha256(args.config.read_bytes()).hexdigest(), + "manifest_sha256": manifest.sha256, + "manifest": [ + { + "name": support.name, + "trajectory_count": len(support.episodes), + "anchor_count": sum(len(episode.anchors) for episode in support.episodes), + } + for support in manifest.datasets + ], + "grid_size": int(args.grid_size), + "interpolation": "linear_within_episode", + "percentile_scope": "within_episode_after_interpolation", + "epsilon": float(args.epsilon), + "alpha": float(args.alpha), + "estimator_seeds": [int(seed) for seed in args.estimator_seeds], + "samples_per_estimator": int(args.samples_per_estimator), + "transform_seed_base": int(args.transform_seed_base), + "model_dtype": args.model_dtype, + } + + +def _initialize_metadata( + path: Path, + *, + identity: dict[str, Any], + episode_count: int, + grid_anchor_count: int, +) -> dict[str, Any]: + if path.exists(): + metadata = json.loads(path.read_text(encoding="utf-8")) + if metadata.get("identity") != identity: + raise ValueError( + "Residual cache metadata belongs to a different scorer protocol" + ) + return metadata + metadata = { + "schema_version": SCHEMA_VERSION, + "status": "in_progress", + "identity": identity, + "planned_episode_count": int(episode_count), + "planned_grid_anchor_count": int(grid_anchor_count), + "completed_episode_count": 0, + "completed_grid_anchor_count": 0, + "scoring_seconds": 0.0, + "proposal": None, + } + _atomic_json(path, metadata) + return metadata + + +def _refresh_progress(metadata: dict[str, Any], cache_dir: Path, manifest) -> None: + completed_episodes = 0 + completed_anchors = 0 + for dataset_index, dataset in enumerate(manifest.datasets): + for episode_index, _ in enumerate(dataset.episodes): + path = _shard_path(cache_dir, dataset_index, episode_index) + if not path.exists(): + continue + payload = json.loads(path.read_text(encoding="utf-8")) + completed_episodes += 1 + completed_anchors += len(payload.get("anchor_positions", [])) + metadata["completed_episode_count"] = completed_episodes + metadata["completed_grid_anchor_count"] = completed_anchors + + +def _finalize_proposal( + *, + manifest, + cache_dir: Path, + identity: dict[str, Any], + output: Path, + alpha: float, + epsilon: float, + estimator_count: int, +) -> TemporalProposal: + masses = [] + for dataset_index, dataset in enumerate(manifest.datasets): + dataset_masses = [] + for episode_index, episode in enumerate(dataset.episodes): + positions = quantile_anchor_positions( + len(episode.anchors), + int(identity["grid_size"]), + ) + base_indices = [int(episode.anchors[int(index)]) for index in positions] + shard = _load_valid_shard( + _shard_path(cache_dir, dataset_index, episode_index), + identity=identity, + dataset_index=dataset_index, + episode_index=episode_index, + trajectory_id=episode.trajectory_id, + positions=positions, + base_indices=base_indices, + estimator_count=estimator_count, + ) + full_scores = interpolate_grid_scores( + len(episode.anchors), + positions, + np.asarray(shard["grid_scores"], dtype=np.float64), + ) + dataset_masses.append( + percentile_intensity_masses(full_scores, epsilon=epsilon) + ) + masses.append(dataset_masses) + + proposal = TemporalProposal.from_masses( + manifest=manifest, + alpha=alpha, + masses=masses, + score_kind="residual", + scorer_sha256="residual-grid-v1", + validation_certificate={ + "protocol": "prospective-grid-residual-v1", + "scorer_identity": identity, + "grid_size": identity["grid_size"], + "interpolation": identity["interpolation"], + "percentile_scope": identity["percentile_scope"], + "epsilon": epsilon, + }, + ) + proposal.save(output) + return proposal + + +def main(argv=None) -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument( + "--manifest", + type=Path, + required=True, + help=( + "Frozen manifest artifact. This explicit argument supports exact " + "common-start configs saved before the generated manifest existed." + ), + ) + parser.add_argument("--cache-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--grid-size", type=int, default=8) + parser.add_argument("--alpha", type=float, default=0.2) + parser.add_argument("--epsilon", type=float, default=0.1) + parser.add_argument( + "--estimator-seeds", + type=int, + nargs="+", + default=[16200, 16201, 16202, 16203], + ) + parser.add_argument("--samples-per-estimator", type=int, default=8) + parser.add_argument("--transform-seed-base", type=int, default=16500) + parser.add_argument( + "--batch-size", + type=int, + default=8, + help="Score this many anchors per VLM/DiT call; K=8 episodes fill one batch.", + ) + parser.add_argument( + "--model-dtype", + choices=("float32", "bfloat16"), + default="float32", + ) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--shard-index", type=int, default=0) + parser.add_argument( + "--max-new-episodes", + type=int, + default=0, + help="Stop cleanly after this many newly scored episodes; 0 means all.", + ) + parser.add_argument("--plan-only", action="store_true") + args = parser.parse_args(argv) + + from starVLA.model.framework.share_tools import apply_config_compat + from starVLA.model.framework.base_framework import build_framework + from starVLA.dataloader.lerobot_datasets import get_vla_dataset + + if args.grid_size <= 0 or args.batch_size <= 0: + raise ValueError("grid-size and batch-size must be positive") + if args.num_shards <= 0: + raise ValueError("num-shards must be positive") + if not 0 <= args.shard_index < args.num_shards: + raise ValueError("shard-index must satisfy 0 <= shard-index < num-shards") + if not 0.0 <= args.alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if args.epsilon <= 0.0: + raise ValueError("epsilon must be positive") + if not args.estimator_seeds: + raise ValueError("at least one estimator seed is required") + if len(set(args.estimator_seeds)) != len(args.estimator_seeds): + raise ValueError("estimator seeds must be unique") + if args.samples_per_estimator <= 0: + raise ValueError("samples-per-estimator must be positive") + + cfg = apply_config_compat(OmegaConf.load(args.config)) + cfg.output_dir = str(args.cache_dir / "runtime") + cfg.datasets.vla_data.temporal_proposal_path = None + configured_manifest = cfg.datasets.vla_data.get("valid_anchor_manifest_path") + requested_manifest = args.manifest.expanduser().resolve(strict=True) + if configured_manifest: + if ( + Path(configured_manifest).expanduser().resolve(strict=True) + != requested_manifest + ): + raise ValueError( + "--manifest differs from config valid_anchor_manifest_path" + ) + else: + cfg.datasets.vla_data.valid_anchor_manifest_path = str( + requested_manifest + ) + dataset = get_vla_dataset( + data_cfg=cfg.datasets.vla_data, + balance_dataset_weights=cfg.datasets.vla_data.get( + "balance_dataset_weights", + False, + ), + balance_trajectory_weights=cfg.datasets.vla_data.get( + "balance_trajectory_weights", + False, + ), + seed=int(cfg.get("seed", 42)), + ) + manifest = dataset.valid_anchor_manifest + if manifest is None: + raise RuntimeError("Residual proposal requires traceable_sampling") + + identity = _build_identity(args, manifest) + episode_count = sum(len(item.episodes) for item in manifest.datasets) + grid_anchor_count = _planned_grid_count(manifest, args.grid_size) + metadata_path = args.cache_dir / "metadata.json" + metadata = _initialize_metadata( + metadata_path, + identity=identity, + episode_count=episode_count, + grid_anchor_count=grid_anchor_count, + ) + _refresh_progress(metadata, args.cache_dir, manifest) + _atomic_json(metadata_path, metadata) + print( + json.dumps( + { + "identity": identity, + "episodes": episode_count, + "grid_anchors": grid_anchor_count, + "completed_episodes": metadata["completed_episode_count"], + "completed_grid_anchors": metadata["completed_grid_anchor_count"], + "estimated_hours_at_0.312_seconds_per_anchor": ( + grid_anchor_count * 0.312 / 3600.0 + ), + }, + indent=2, + ), + flush=True, + ) + if args.plan_only: + return + + parallel_worker = args.num_shards > 1 + assigned_episode_count = 0 + assigned_grid_anchor_count = 0 + global_episode_index = 0 + for dataset_support in manifest.datasets: + for episode in dataset_support.episodes: + if _episode_belongs_to_shard( + global_episode_index, + shard_index=args.shard_index, + num_shards=args.num_shards, + ): + assigned_episode_count += 1 + assigned_grid_anchor_count += min( + len(episode.anchors), + args.grid_size, + ) + global_episode_index += 1 + + missing_exists = metadata["completed_episode_count"] < episode_count + if missing_exists: + model = build_framework(cfg) + checkpoint = torch.load( + args.checkpoint, + map_location="cpu", + weights_only=True, + ) + model.load_state_dict(checkpoint, strict=True) + model.requires_grad_(False).eval().to( + device="cuda", + dtype=getattr(torch, args.model_dtype), + ) + else: + model = None + + new_episodes = 0 + completed_assigned_episodes = 0 + completed_assigned_anchors = 0 + session_scoring_seconds = 0.0 + worker_progress = { + "schema_version": 1, + "status": "in_progress", + "identity": identity, + "shard_index": args.shard_index, + "num_shards": args.num_shards, + "assigned_episode_count": assigned_episode_count, + "assigned_grid_anchor_count": assigned_grid_anchor_count, + "completed_episode_count": 0, + "completed_grid_anchor_count": 0, + "new_episode_count": 0, + "scoring_seconds": 0.0, + } + worker_path = _worker_progress_path( + args.cache_dir, + args.shard_index, + args.num_shards, + ) + if parallel_worker: + _atomic_json(worker_path, worker_progress) + try: + global_episode_index = 0 + for dataset_index, dataset_support in enumerate(manifest.datasets): + source_dataset = dataset.datasets[dataset_index] + for episode_index, episode in enumerate(dataset_support.episodes): + selected_for_worker = _episode_belongs_to_shard( + global_episode_index, + shard_index=args.shard_index, + num_shards=args.num_shards, + ) + global_episode_index += 1 + if not selected_for_worker: + continue + shard_path = _shard_path( + args.cache_dir, + dataset_index, + episode_index, + ) + positions = quantile_anchor_positions( + len(episode.anchors), + args.grid_size, + ) + base_indices = [ + int(episode.anchors[int(index)]) for index in positions + ] + if shard_path.exists(): + _load_valid_shard( + shard_path, + identity=identity, + dataset_index=dataset_index, + episode_index=episode_index, + trajectory_id=episode.trajectory_id, + positions=positions, + base_indices=base_indices, + estimator_count=len(args.estimator_seeds), + ) + completed_assigned_episodes += 1 + completed_assigned_anchors += len(positions) + continue + + if model is None: + raise RuntimeError("Missing shard discovered after model skip") + examples, loaded_base_indices, transform_seeds = _load_examples( + source_dataset, + episode, + positions, + dataset_index=dataset_index, + episode_index=episode_index, + transform_seed_base=args.transform_seed_base, + ) + if loaded_base_indices != base_indices: + raise RuntimeError("Loaded base indices differ from frozen plan") + started = time.perf_counter() + estimators = _score_examples( + model, + examples, + estimator_seeds=args.estimator_seeds, + samples_per_estimator=args.samples_per_estimator, + batch_size=args.batch_size, + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + session_scoring_seconds += elapsed + grid_scores = estimators.mean(axis=1) + shard = { + "schema_version": SCHEMA_VERSION, + "identity": identity, + "dataset_index": dataset_index, + "dataset_name": dataset_support.name, + "episode_index": episode_index, + "trajectory_id": _json_scalar(episode.trajectory_id), + "anchor_count": len(episode.anchors), + "anchor_positions": positions.tolist(), + "base_indices": base_indices, + "transform_seeds": transform_seeds, + "estimator_residuals": estimators.tolist(), + "grid_scores": grid_scores.tolist(), + "scoring_seconds": elapsed, + } + _atomic_json(shard_path, shard) + new_episodes += 1 + completed_assigned_episodes += 1 + completed_assigned_anchors += len(positions) + if parallel_worker: + worker_progress.update( + { + "completed_episode_count": completed_assigned_episodes, + "completed_grid_anchor_count": completed_assigned_anchors, + "new_episode_count": new_episodes, + "scoring_seconds": session_scoring_seconds, + } + ) + if new_episodes % 25 == 0: + _atomic_json(worker_path, worker_progress) + print( + f"residual-cache shard={args.shard_index}/" + f"{args.num_shards} episode=" + f"{completed_assigned_episodes}/" + f"{assigned_episode_count} seconds={elapsed:.3f}", + flush=True, + ) + else: + metadata["completed_episode_count"] += 1 + metadata["completed_grid_anchor_count"] += len(positions) + metadata["scoring_seconds"] += elapsed + _atomic_json(metadata_path, metadata) + print( + f"residual-cache episode={metadata['completed_episode_count']}/" + f"{episode_count} grid_anchors={len(positions)} " + f"seconds={elapsed:.3f}", + flush=True, + ) + if ( + args.max_new_episodes > 0 + and new_episodes >= args.max_new_episodes + ): + if parallel_worker: + worker_progress["status"] = "partial" + _atomic_json(worker_path, worker_progress) + print("Stopped at max-new-episodes; cache is resumable.", flush=True) + return + finally: + if model is not None: + del model + torch.cuda.empty_cache() + + if parallel_worker: + worker_progress.update( + { + "status": "complete", + "completed_episode_count": completed_assigned_episodes, + "completed_grid_anchor_count": completed_assigned_anchors, + "new_episode_count": new_episodes, + "scoring_seconds": session_scoring_seconds, + } + ) + _atomic_json(worker_path, worker_progress) + if completed_assigned_episodes != assigned_episode_count: + raise RuntimeError("Residual cache worker finished with missing episodes") + print(json.dumps(worker_progress, indent=2, sort_keys=True), flush=True) + return + + _refresh_progress(metadata, args.cache_dir, manifest) + if metadata["completed_episode_count"] != episode_count: + metadata["status"] = "in_progress" + _atomic_json(metadata_path, metadata) + raise RuntimeError("Residual cache is incomplete after scoring loop") + + proposal = _finalize_proposal( + manifest=manifest, + cache_dir=args.cache_dir, + identity=identity, + output=args.output, + alpha=args.alpha, + epsilon=args.epsilon, + estimator_count=len(args.estimator_seeds), + ) + metadata["status"] = "complete" + metadata["proposal"] = { + "path": str(args.output.resolve()), + "sha256": proposal.sha256, + "importance_weight_min": min( + float(np.min(weights)) + for dataset_weights in proposal.importance_weights + for weights in dataset_weights + ), + "importance_weight_max": max( + float(np.max(weights)) + for dataset_weights in proposal.importance_weights + for weights in dataset_weights + ), + } + _atomic_json(metadata_path, metadata) + print( + json.dumps( + { + "status": "complete", + "proposal": metadata["proposal"], + "scoring_seconds": metadata["scoring_seconds"], + "session_scoring_seconds": session_scoring_seconds, + }, + indent=2, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/real_robot/source/retime/commands/shuffle.py b/real_robot/source/retime/commands/shuffle.py new file mode 100644 index 0000000000000000000000000000000000000000..b59721f6a6fdaa4905f84f129a4f1f0f607c6183 --- /dev/null +++ b/real_robot/source/retime/commands/shuffle.py @@ -0,0 +1,47 @@ +"""Build the Shuffle control by permuting residual masses per trajectory.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from retime.sampler import ValidAnchorManifest +from retime.proposal import TemporalProposal +from retime.scoring.shuffle import build_shuffle_proposal + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--residual-proposal", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--permutation-seed", type=int, required=True) + args = parser.parse_args() + + manifest = ValidAnchorManifest.load(args.manifest) + residual = TemporalProposal.load(args.residual_proposal, manifest) + shuffled = build_shuffle_proposal( + manifest=manifest, + residual_proposal=residual, + permutation_seed=args.permutation_seed, + ) + shuffled.save(args.output) + print( + json.dumps( + { + "status": "complete", + "manifest_sha256": manifest.sha256, + "source_proposal_sha256": residual.sha256, + "shuffle_proposal_sha256": shuffled.sha256, + "permutation_seed": args.permutation_seed, + "output": str(args.output.resolve()), + }, + indent=2, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/real_robot/source/retime/objective.py b/real_robot/source/retime/objective.py new file mode 100644 index 0000000000000000000000000000000000000000..52ed7bce20458c982f76e9b414162b812505e9e2 --- /dev/null +++ b/real_robot/source/retime/objective.py @@ -0,0 +1,152 @@ +"""Objective-preserving primitives for InfoFlow training. + +This module intentionally depends only on PyTorch. It contains no dataset or +model-specific policy, so its contracts can be tested before sampler wiring. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class SafeMixture: + """A normalized safe proposal and its exact importance correction.""" + + base_probability: torch.Tensor + proposal_probability: torch.Tensor + sampling_probability: torch.Tensor + importance_weight: torch.Tensor + max_importance_weight: float + + +def _normalized_probability(probability: torch.Tensor, *, name: str) -> torch.Tensor: + if not isinstance(probability, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if probability.ndim != 1 or probability.numel() == 0: + raise ValueError(f"{name} must be a non-empty one-dimensional tensor") + if probability.dtype == torch.bool or not probability.dtype.is_floating_point: + raise TypeError(f"{name} must have a floating-point dtype") + + probability64 = probability.detach().to(dtype=torch.float64) + if not torch.isfinite(probability64).all(): + raise ValueError(f"{name} must contain only finite values") + if (probability64 < 0).any(): + raise ValueError(f"{name} must be non-negative") + + total = probability64.sum() + if total <= 0: + raise ValueError(f"{name} must have a positive sum") + return probability64 / total + + +def build_safe_mixture( + base_probability: torch.Tensor, + proposal_probability: torch.Tensor, + *, + alpha: float, +) -> SafeMixture: + """Build ``q = (1 - alpha) p + alpha r`` and exact ``p / q`` weights.""" + + if base_probability.device != proposal_probability.device: + raise ValueError("base_probability and proposal_probability must be on the same device") + if base_probability.shape != proposal_probability.shape: + raise ValueError("base_probability and proposal_probability must have the same shape") + if not isinstance(alpha, (int, float)) or isinstance(alpha, bool) or not math.isfinite(float(alpha)): + raise TypeError("alpha must be a finite real number") + alpha = float(alpha) + if not 0.0 <= alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + + p = _normalized_probability(base_probability, name="base_probability") + r = _normalized_probability(proposal_probability, name="proposal_probability") + q = (1.0 - alpha) * p + alpha * r + + positive_support = p > 0 + if not (q[positive_support] > 0).all(): + raise RuntimeError("safe mixture lost support for a positive-probability baseline item") + + weight = torch.zeros_like(p) + weight[positive_support] = p[positive_support] / q[positive_support] + max_weight = 1.0 / (1.0 - alpha) + if (weight > max_weight + 1e-12).any(): + raise RuntimeError("safe-mixture importance bound was violated") + + return SafeMixture( + base_probability=p, + proposal_probability=r, + sampling_probability=q, + importance_weight=weight.detach(), + max_importance_weight=max_weight, + ) + + +def reduce_flow_matching_loss(elementwise_loss: torch.Tensor, *, reduction: str = "mean") -> torch.Tensor: + """Reduce flow loss while retaining the original batch dimension.""" + + if not isinstance(elementwise_loss, torch.Tensor): + raise TypeError("elementwise_loss must be a torch.Tensor") + if elementwise_loss.ndim == 0 or elementwise_loss.shape[0] == 0: + raise ValueError("elementwise_loss must have a non-empty batch dimension") + if reduction not in {"mean", "none"}: + raise ValueError("reduction must be either 'mean' or 'none'") + + # Preserve the legacy scalar path's reduction order exactly. Computing + # per-sample means first can introduce a tiny floating-point difference. + if reduction == "mean": + return elementwise_loss.mean() + + if elementwise_loss.ndim == 1: + per_sample = elementwise_loss + else: + per_sample = elementwise_loss.mean(dim=tuple(range(1, elementwise_loss.ndim))) + return per_sample + + +def aggregate_repeated_losses( + repeated_per_sample_loss: torch.Tensor, + *, + repeat_count: int, + batch_size: int | None = None, +) -> torch.Tensor: + """Average repeat-major stochastic flow draws back to shape ``[B]``.""" + + if not isinstance(repeated_per_sample_loss, torch.Tensor): + raise TypeError("repeated_per_sample_loss must be a torch.Tensor") + if repeated_per_sample_loss.ndim != 1 or repeated_per_sample_loss.numel() == 0: + raise ValueError("repeated_per_sample_loss must be a non-empty one-dimensional tensor") + if not isinstance(repeat_count, int) or isinstance(repeat_count, bool) or repeat_count <= 0: + raise ValueError("repeat_count must be a positive integer") + if repeated_per_sample_loss.numel() % repeat_count != 0: + raise ValueError("loss count must be divisible by repeat_count") + + inferred_batch_size = repeated_per_sample_loss.numel() // repeat_count + if batch_size is not None and batch_size != inferred_batch_size: + raise ValueError("batch_size does not match the repeated loss shape") + return repeated_per_sample_loss.reshape(repeat_count, inferred_batch_size).mean(dim=0) + + +def ordinary_importance_corrected_mean( + per_sample_loss: torch.Tensor, + importance_weight: torch.Tensor, +) -> torch.Tensor: + """Return the ordinary (not self-normalized) importance estimator.""" + + if not isinstance(per_sample_loss, torch.Tensor) or not isinstance(importance_weight, torch.Tensor): + raise TypeError("per_sample_loss and importance_weight must be torch.Tensor instances") + if per_sample_loss.ndim != 1 or per_sample_loss.numel() == 0: + raise ValueError("per_sample_loss must be a non-empty one-dimensional tensor") + if importance_weight.shape != per_sample_loss.shape: + raise ValueError("importance_weight must have the same shape as per_sample_loss") + + detached_weight = importance_weight.detach() + if not torch.isfinite(detached_weight).all() or (detached_weight < 0).any(): + raise ValueError("importance_weight must be finite and non-negative") + + work_dtype = torch.float64 if per_sample_loss.dtype == torch.float64 else torch.float32 + loss = per_sample_loss.to(dtype=work_dtype) + weight = detached_weight.to(device=per_sample_loss.device, dtype=work_dtype) + return (weight * loss).mean() diff --git a/real_robot/source/retime/prepare.py b/real_robot/source/retime/prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..7809e28b9a423039d46e82901df0f7504f26b430 --- /dev/null +++ b/real_robot/source/retime/prepare.py @@ -0,0 +1,32 @@ +"""Derive temporal anchor support directly from the configured dataset.""" + +import argparse +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + from omegaconf import OmegaConf + from starVLA.dataloader.lerobot_datasets import get_vla_dataset + from starVLA.model.framework.share_tools import apply_config_compat + + cfg = apply_config_compat(OmegaConf.load(args.config)) + data = cfg.datasets.vla_data + data.temporal_proposal_path = None + data.valid_anchor_manifest_path = None + data.traceable_sampling = True + dataset = get_vla_dataset( + data_cfg=data, + balance_dataset_weights=data.get("balance_dataset_weights", False), + balance_trajectory_weights=data.get("balance_trajectory_weights", False), + seed=int(cfg.get("seed", 42)), + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + dataset.valid_anchor_manifest.save(args.output) + + +if __name__ == "__main__": + main() diff --git a/real_robot/source/retime/pretrain.py b/real_robot/source/retime/pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..23a7929df686fb1da64d7f515800ff05c4edafb9 --- /dev/null +++ b/real_robot/source/retime/pretrain.py @@ -0,0 +1,200 @@ +"""Train a reviewer-owned common-start policy and save exact resume state.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import subprocess +import sys + + +def parser() -> argparse.ArgumentParser: + result = argparse.ArgumentParser(description=__doc__) + result.add_argument("--config", type=Path, required=True) + result.add_argument("--accelerate-config", type=Path, required=True) + result.add_argument("--base-steps", type=int, required=True) + result.add_argument("--vlm-path", type=Path, required=True) + result.add_argument("--data-root", type=Path, required=True) + result.add_argument("--data-mix", type=str) + result.add_argument("--statistics-path", type=Path) + result.add_argument("--output", type=Path, required=True) + result.add_argument("--enable-wandb", action="store_true") + result.add_argument( + "--omit-data-identity-files", + action="store_true", + help="Record computed dataset identity in the manifest without separate identity files.", + ) + result.add_argument( + "--backend-root", + type=Path, + default=Path(__file__).resolve().parents[1], + ) + result.add_argument("--dry-run", action="store_true") + return result + + +def _file(path: Path, label: str) -> Path: + path = path.expanduser().resolve() + if not path.is_file(): + raise ValueError(f"{label} is not a file: {path}") + return path + + +def _directory(path: Path, label: str) -> Path: + path = path.expanduser().resolve() + if not path.is_dir(): + raise ValueError(f"{label} is not a directory: {path}") + return path + + +def build_plan(args: argparse.Namespace) -> dict: + from omegaconf import OmegaConf + + if args.base_steps < 1: + raise ValueError("--base-steps must be positive") + root = args.backend_root.expanduser().resolve() + source_config = _file(args.config, "config") + launch_config = _file(args.accelerate_config, "accelerate config") + training = _file( + root / "starVLA" / "training" / "train_starvla.py", + "training backend", + ) + cfg = OmegaConf.load(source_config) + if not args.enable_wandb: + cfg.wandb_project = None + launch = OmegaConf.load(launch_config) + reported_stop = int(cfg.trainer.stop_after_steps) + if reported_stop - args.base_steps != 500: + raise ValueError( + "--base-steps must be exactly 500 updates before the selected continuation stop" + ) + world_size = int(launch.num_processes) + if int(OmegaConf.select(launch, "num_machines", default=1)) != 1: + raise ValueError("The reviewer pretraining launcher supports one machine") + if world_size != int(cfg.datasets.vla_data.expected_world_size): + raise ValueError("Accelerate process count differs from the dataset training contract") + accumulation = int(cfg.trainer.gradient_accumulation_steps) + ds_accumulation = OmegaConf.select( + launch, + "deepspeed_config.gradient_accumulation_steps", + ) + if ds_accumulation is not None and int(ds_accumulation) != accumulation: + raise ValueError("Accelerate and trainer gradient accumulation differ") + + output = args.output.expanduser().resolve() + if output.exists() and any(output.iterdir()): + raise ValueError(f"Pretraining output is not empty: {output}") + cfg.run_root_dir = str(output.parent) + cfg.run_id = output.name + cfg.output_dir = str(output) + cfg.framework.qwenvl.base_vlm = str(_directory(args.vlm_path, "VLM backbone")) + cfg.datasets.vla_data.data_root_dir = str(_directory(args.data_root, "dataset root")) + if args.data_mix: + cfg.datasets.vla_data.data_mix = args.data_mix + if args.statistics_path: + cfg.datasets.vla_data.precomputed_statistics_path = str( + _file(args.statistics_path, "statistics") + ) + if args.omit_data_identity_files: + cfg.datasets.vla_data.dataset_provenance_path = None + cfg.datasets.vla_data.dataset_content_fingerprint_path = None + + manifest = output / "manifest.json" + cfg.datasets.vla_data.valid_anchor_manifest_path = str(manifest) + cfg.datasets.vla_data.temporal_proposal_path = None + cfg.datasets.vla_data.transform_rng_seed = None + cfg.seed = int(OmegaConf.select(cfg, "trainer.branch_parent_seed", default=42)) + cfg.trainer.is_resume = False + cfg.trainer.allow_model_only_resume = False + cfg.trainer.pretrained_checkpoint = None + cfg.trainer.resume_accelerate_checkpoint = None + cfg.trainer.branch_parent_seed = None + cfg.trainer.branch_rng_seed = None + cfg.trainer.model_rng_seed = None + cfg.trainer.stop_after_steps = args.base_steps + cfg.trainer.save_steps = [args.base_steps] + cfg.trainer.save_final_checkpoint = False + cfg.trainer.save_consolidated_model = True + cfg.trainer.evaluation_snapshot_interval = None + + materialized_config = output / "config.yaml" + config_text = OmegaConf.to_yaml(cfg, resolve=True) + commands = [ + [ + sys.executable, + "-m", + "retime.prepare", + "--config", + str(materialized_config), + "--output", + str(manifest), + ], + [ + sys.executable, + "-m", + "accelerate.commands.launch", + "--config_file", + str(launch_config), + "--num_processes", + str(world_size), + str(training), + "--config_yaml", + str(materialized_config), + ], + ] + return { + "output": str(output), + "base_steps": args.base_steps, + "manifest": str(manifest), + "global_batch_size": ( + int(cfg.datasets.vla_data.per_device_batch_size) + * world_size + * accumulation + ), + "config_text": config_text, + "commands": commands, + } + + +def main(argv: list[str] | None = None) -> int: + arguments = parser() + args = arguments.parse_args(argv) + try: + plan = build_plan(args) + except (ValueError, OSError, KeyError, ImportError) as error: + arguments.error(str(error)) + print( + json.dumps( + {key: value for key, value in plan.items() if key != "config_text"}, + indent=2, + ), + flush=True, + ) + if args.dry_run: + return 0 + output = Path(plan["output"]) + output.mkdir(parents=True, exist_ok=True) + (output / "config.yaml").write_text(plan["config_text"], encoding="utf-8") + (output / "pretrain-plan.json").write_text( + json.dumps( + {key: value for key, value in plan.items() if key != "config_text"}, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + environment = dict(os.environ) + environment["PYTHONPATH"] = str(args.backend_root.expanduser().resolve()) + ( + os.pathsep + environment["PYTHONPATH"] + if environment.get("PYTHONPATH") + else "" + ) + for command in plan["commands"]: + subprocess.run(command, cwd=args.backend_root.expanduser().resolve(), env=environment, check=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/real_robot/source/retime/proposal.py b/real_robot/source/retime/proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..be5afa058b38a1d825a3ec56d8f3bb02c57371fd --- /dev/null +++ b/real_robot/source/retime/proposal.py @@ -0,0 +1,223 @@ +"""Offline temporal proposals that preserve StarVLA's dataset/episode marginals.""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Sequence + +import numpy as np + +from retime.sampler import ( + LogicalDraw, + SampleTrace, + ValidAnchorManifest, + _stateless_seed, +) + + +class TemporalProposal: + schema_version = 1 + valid_score_kinds = { + "synthetic", + "random", + "shuffle", + "action_change", + "residual", + "likelihood", + } + + def __init__( + self, + manifest: ValidAnchorManifest, + alpha: float, + anchor_probabilities: Sequence[Sequence[Sequence[float]]], + score_kind: str, + scorer_sha256: str, + validation_certificate: dict[str, Any] | None = None, + ) -> None: + if not 0.0 <= alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if score_kind not in self.valid_score_kinds: + raise ValueError(f"unsupported score_kind: {score_kind}") + if len(anchor_probabilities) != len(manifest.datasets): + raise ValueError("proposal dataset count does not match the manifest") + + normalized_probabilities = [] + importance_weights = [] + normalized_masses = [] + for dataset_index, (dataset, dataset_probabilities) in enumerate( + zip(manifest.datasets, anchor_probabilities) + ): + if len(dataset_probabilities) != len(dataset.episodes): + raise ValueError( + f"proposal episode count differs for dataset {dataset_index}" + ) + normalized_dataset = [] + dataset_weights = [] + normalized_dataset_masses = [] + for episode_index, (episode, proposal_mass) in enumerate( + zip(dataset.episodes, dataset_probabilities) + ): + mass = np.asarray(proposal_mass, dtype=np.float64) + if mass.shape != (len(episode.anchors),): + raise ValueError( + f"proposal anchor count differs for dataset {dataset_index}, " + f"episode {episode_index}" + ) + if not np.all(np.isfinite(mass)) or np.any(mass < 0) or mass.sum() <= 0: + raise ValueError("proposal masses must be finite, non-negative, and non-zero") + residual_probability = mass / mass.sum(dtype=np.float64) + baseline_probability = np.full(len(mass), 1.0 / len(mass), dtype=np.float64) + safe_probability = (1.0 - alpha) * baseline_probability + alpha * residual_probability + weights = baseline_probability / safe_probability + if not np.isclose(safe_probability.sum(dtype=np.float64), 1.0, atol=1e-12): + raise ValueError("safe temporal proposal does not sum to one") + if np.max(weights) > 1.0 / (1.0 - alpha) + 1e-12: + raise ValueError("importance weight exceeds the analytical safe-mixture bound") + normalized_dataset.append(safe_probability) + dataset_weights.append(weights) + normalized_dataset_masses.append(mass.copy()) + normalized_probabilities.append(tuple(normalized_dataset)) + importance_weights.append(tuple(dataset_weights)) + normalized_masses.append(tuple(normalized_dataset_masses)) + + self.manifest_sha256 = manifest.sha256 + self.manifest_support = manifest.to_dict(include_hash=False)["datasets"] + self.alpha = float(alpha) + self.score_kind = score_kind + self.scorer_sha256 = scorer_sha256 + self.validation_certificate = validation_certificate + self.proposal_masses = tuple(normalized_masses) + self.anchor_probabilities = tuple(normalized_probabilities) + self.importance_weights = tuple(importance_weights) + canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":")) + self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + @classmethod + def from_masses( + cls, + manifest: ValidAnchorManifest, + alpha: float, + masses: Sequence[Sequence[Sequence[float]]], + score_kind: str = "synthetic", + scorer_sha256: str = "synthetic-fixture", + validation_certificate: dict[str, Any] | None = None, + ) -> "TemporalProposal": + return cls( + manifest=manifest, + alpha=alpha, + anchor_probabilities=masses, + score_kind=score_kind, + scorer_sha256=scorer_sha256, + validation_certificate=validation_certificate, + ) + + @classmethod + def from_dict( + cls, manifest: ValidAnchorManifest, payload: dict[str, Any] + ) -> "TemporalProposal": + if payload.get("schema_version") != cls.schema_version: + raise ValueError( + f"unsupported temporal proposal schema: {payload.get('schema_version')}" + ) + expected_support = payload.get("manifest_support") + current_support = json.loads( + json.dumps(manifest.to_dict(include_hash=False)["datasets"]) + ) + if expected_support is not None and expected_support != current_support: + raise ValueError("temporal proposal changes dataset/episode/anchor support") + proposal = cls( + manifest=manifest, + alpha=float(payload["alpha"]), + anchor_probabilities=payload["proposal_masses"], + score_kind=str(payload["score_kind"]), + scorer_sha256=str(payload["scorer_sha256"]), + validation_certificate=payload.get("validation_certificate"), + ) + return proposal + + @classmethod + def load( + cls, path: str | Path, manifest: ValidAnchorManifest + ) -> "TemporalProposal": + with open(path, "r", encoding="utf-8") as proposal_file: + return cls.from_dict(manifest, json.load(proposal_file)) + + def save(self, path: str | Path) -> None: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".tmp") + with open(temporary, "w", encoding="utf-8") as proposal_file: + json.dump(self.to_dict(), proposal_file, indent=2) + os.replace(temporary, destination) + + def sample( + self, manifest: ValidAnchorManifest, seed: int, draw: LogicalDraw + ) -> SampleTrace: + rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id)) + dataset_index = int(rng.choice(len(manifest.datasets), p=manifest._dataset_probabilities)) + dataset = manifest.datasets[dataset_index] + trajectory_index = int( + rng.choice( + len(dataset.episodes), p=manifest._episode_probabilities[dataset_index] + ) + ) + episode = dataset.episodes[trajectory_index] + if self.alpha == 0.0: + # Preserve the exact baseline RNG path, not merely its distribution. + anchor_index = int(rng.integers(len(episode.anchors))) + else: + anchor_index = int( + rng.choice( + len(episode.anchors), + p=self.anchor_probabilities[dataset_index][trajectory_index], + ) + ) + baseline_probability = ( + dataset.probability + * dataset.episode_probabilities[trajectory_index] + / len(episode.anchors) + ) + proposal_probability = ( + dataset.probability + * dataset.episode_probabilities[trajectory_index] + * self.anchor_probabilities[dataset_index][trajectory_index][anchor_index] + ) + return SampleTrace( + manifest_sha256=manifest.sha256, + epoch=draw.epoch, + global_draw_id=draw.global_draw_id, + dataset_index=dataset_index, + dataset_name=dataset.name, + trajectory_index=episode.trajectory_index, + trajectory_id=episode.trajectory_id, + anchor_index=anchor_index, + base_index=episode.anchors[anchor_index], + baseline_probability=float(baseline_probability), + proposal_sha256=self.sha256, + proposal_probability=float(proposal_probability), + importance_weight=float( + self.importance_weights[dataset_index][trajectory_index][anchor_index] + ), + ) + + def to_dict(self, include_hash: bool = True) -> dict[str, Any]: + payload = { + "schema_version": self.schema_version, + "manifest_sha256": self.manifest_sha256, + "manifest_support": self.manifest_support, + "alpha": self.alpha, + "score_kind": self.score_kind, + "scorer_sha256": self.scorer_sha256, + "validation_certificate": self.validation_certificate, + "proposal_masses": [ + [probabilities.tolist() for probabilities in dataset] + for dataset in self.proposal_masses + ], + } + if include_hash: + payload["sha256"] = self.sha256 + return payload diff --git a/real_robot/source/retime/sampler.py b/real_robot/source/retime/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..fbe3f06c53babf5c34661686838c4ab746c0d38f --- /dev/null +++ b/real_robot/source/retime/sampler.py @@ -0,0 +1,332 @@ +"""Deterministic hierarchical sampling primitives for InfoFlow experiments.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Iterator, Sequence + +import numpy as np +from torch.utils.data import Sampler + + +def _identity_files(dataset_root: Path) -> list[Path]: + files = list((dataset_root / "data").rglob("*.parquet")) + files.extend((dataset_root / "videos").rglob("*.mp4")) + meta_root = dataset_root / "meta" + if meta_root.exists(): + for pattern in ("*.json", "*.jsonl", "*.parquet"): + files.extend(meta_root.rglob(pattern)) + return sorted(set(files), key=lambda path: path.relative_to(dataset_root).as_posix()) + + +def compute_dataset_content_fingerprint(dataset_root: str | Path) -> dict[str, Any]: + """Record the immutable LeRobot data and metadata file inventory.""" + root = Path(dataset_root) + entries = [] + total_bytes = 0 + for file_path in _identity_files(root): + relative_path = file_path.relative_to(root).as_posix() + file_size = file_path.stat().st_size + entries.append({"path": relative_path, "bytes": file_size}) + total_bytes += file_size + return { + "file_count": len(entries), + "total_bytes": total_bytes, + "files": entries, + } + + +def validate_dataset_content_inventory( + dataset_root: str | Path, fingerprint: dict[str, Any] +) -> None: + """Check that the current file inventory matches the recorded dataset inventory.""" + root = Path(dataset_root) + current_entries = [ + { + "path": path.relative_to(root).as_posix(), + "bytes": path.stat().st_size, + } + for path in _identity_files(root) + ] + if current_entries != fingerprint.get("files"): + raise RuntimeError( + f"Dataset file inventory differs from its content fingerprint: {root}" + ) + if len(current_entries) != fingerprint.get("file_count"): + raise RuntimeError("Dataset content fingerprint has an inconsistent file count") + if sum(entry["bytes"] for entry in current_entries) != fingerprint.get("total_bytes"): + raise RuntimeError("Dataset content fingerprint has an inconsistent byte count") + + +def _normalize_probabilities(values: Sequence[float], name: str) -> tuple[float, ...]: + probabilities = np.asarray(values, dtype=np.float64) + if probabilities.ndim != 1 or probabilities.size == 0: + raise ValueError(f"{name} must be a non-empty one-dimensional sequence") + if not np.all(np.isfinite(probabilities)) or np.any(probabilities <= 0): + raise ValueError(f"{name} must contain finite, strictly positive values") + probabilities /= probabilities.sum(dtype=np.float64) + return tuple(float(value) for value in probabilities) + + +def _stateless_seed(seed: int, epoch: int, global_draw_id: int) -> int: + payload = f"{int(seed)}:{int(epoch)}:{int(global_draw_id)}".encode("utf-8") + return int.from_bytes(hashlib.sha256(payload).digest()[:8], "little", signed=False) + + +@dataclass(frozen=True) +class EpisodeSupport: + trajectory_index: int + trajectory_id: int | str + anchors: tuple[int, ...] + + def __post_init__(self) -> None: + if not self.anchors: + raise ValueError("episode support must contain at least one anchor") + if len(set(self.anchors)) != len(self.anchors): + raise ValueError("episode anchors must be unique") + + +@dataclass(frozen=True) +class DatasetSupport: + name: str + probability: float + episodes: tuple[EpisodeSupport, ...] + episode_probabilities: tuple[float, ...] + + def __post_init__(self) -> None: + if not self.episodes: + raise ValueError(f"dataset {self.name!r} has no valid episodes") + if len(self.episodes) != len(self.episode_probabilities): + raise ValueError("episode support and probability lengths differ") + normalized = _normalize_probabilities(self.episode_probabilities, "episode_probabilities") + if not np.allclose(normalized, self.episode_probabilities, rtol=0, atol=1e-12): + raise ValueError("episode probabilities must already sum to one") + + +@dataclass(frozen=True) +class LogicalDraw: + epoch: int + global_draw_id: int + + def __post_init__(self) -> None: + if self.epoch < 0 or self.global_draw_id < 0: + raise ValueError("epoch and global_draw_id must be non-negative") + + +@dataclass(frozen=True) +class SampleTrace: + manifest_sha256: str + epoch: int + global_draw_id: int + dataset_index: int + dataset_name: str + trajectory_index: int + trajectory_id: int | str + anchor_index: int + base_index: int + baseline_probability: float + proposal_sha256: str | None = None + proposal_probability: float | None = None + importance_weight: float = 1.0 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +class ValidAnchorManifest: + schema_version = 1 + + def __init__(self, datasets: Sequence[DatasetSupport], identity: dict[str, Any] | None = None): + if not datasets: + raise ValueError("manifest must contain at least one dataset") + dataset_probabilities = _normalize_probabilities( + [dataset.probability for dataset in datasets], "dataset_probabilities" + ) + if not np.allclose( + dataset_probabilities, + [dataset.probability for dataset in datasets], + rtol=0, + atol=1e-12, + ): + raise ValueError("dataset probabilities must already sum to one") + self.datasets = tuple(datasets) + self.identity = identity or {} + self._dataset_probabilities = np.asarray( + [dataset.probability for dataset in self.datasets], dtype=np.float64 + ) + self._episode_probabilities = tuple( + np.asarray(dataset.episode_probabilities, dtype=np.float64) + for dataset in self.datasets + ) + canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":")) + self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + @classmethod + def from_hierarchy( + cls, + dataset_names: Sequence[str], + dataset_weights: Sequence[float], + episodes: Sequence[Sequence[tuple[int | str, Sequence[int], float]]], + identity: dict[str, Any] | None = None, + ) -> "ValidAnchorManifest": + if len(dataset_names) != len(dataset_weights) or len(dataset_names) != len(episodes): + raise ValueError("dataset names, weights, and episode supports must have equal lengths") + dataset_probabilities = _normalize_probabilities(dataset_weights, "dataset_weights") + supports = [] + for dataset_index, (name, dataset_episodes) in enumerate(zip(dataset_names, episodes)): + if not dataset_episodes: + raise ValueError(f"dataset {name!r} has no valid episodes") + episode_probabilities = _normalize_probabilities( + [entry[2] for entry in dataset_episodes], f"episode_weights[{dataset_index}]" + ) + episode_supports = tuple( + EpisodeSupport( + trajectory_index=trajectory_index, + trajectory_id=trajectory_id, + anchors=tuple(int(anchor) for anchor in anchors), + ) + for trajectory_index, (trajectory_id, anchors, _) in enumerate(dataset_episodes) + ) + supports.append( + DatasetSupport( + name=str(name), + probability=dataset_probabilities[dataset_index], + episodes=episode_supports, + episode_probabilities=episode_probabilities, + ) + ) + return cls(supports, identity=identity) + + @classmethod + def from_dict(cls, payload: dict[str, Any]) -> "ValidAnchorManifest": + datasets = [] + for dataset_payload in payload["datasets"]: + episodes = tuple( + EpisodeSupport( + trajectory_index=int(episode["trajectory_index"]), + trajectory_id=episode["trajectory_id"], + anchors=tuple(int(anchor) for anchor in episode["anchors"]), + ) + for episode in dataset_payload["episodes"] + ) + datasets.append( + DatasetSupport( + name=str(dataset_payload["name"]), + probability=float(dataset_payload["probability"]), + episodes=episodes, + episode_probabilities=tuple( + float(value) for value in dataset_payload["episode_probabilities"] + ), + ) + ) + return cls(datasets, identity=payload.get("identity", {})) + + @classmethod + def load(cls, path: str | Path) -> "ValidAnchorManifest": + with open(path, "r", encoding="utf-8") as manifest_file: + return cls.from_dict(json.load(manifest_file)) + + def save(self, path: str | Path) -> None: + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".tmp") + with open(temporary, "w", encoding="utf-8") as manifest_file: + json.dump(self.to_dict(), manifest_file, indent=2) + os.replace(temporary, destination) + + def sample(self, seed: int, draw: LogicalDraw) -> SampleTrace: + rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id)) + dataset_index = int(rng.choice(len(self.datasets), p=self._dataset_probabilities)) + dataset = self.datasets[dataset_index] + trajectory_index = int( + rng.choice(len(dataset.episodes), p=self._episode_probabilities[dataset_index]) + ) + episode = dataset.episodes[trajectory_index] + anchor_index = int(rng.integers(len(episode.anchors))) + baseline_probability = ( + dataset.probability + * dataset.episode_probabilities[trajectory_index] + / len(episode.anchors) + ) + return SampleTrace( + manifest_sha256=self.sha256, + epoch=int(draw.epoch), + global_draw_id=int(draw.global_draw_id), + dataset_index=dataset_index, + dataset_name=dataset.name, + trajectory_index=episode.trajectory_index, + trajectory_id=episode.trajectory_id, + anchor_index=anchor_index, + base_index=episode.anchors[anchor_index], + baseline_probability=float(baseline_probability), + ) + + def to_dict(self, include_hash: bool = True) -> dict[str, Any]: + payload = { + "schema_version": self.schema_version, + "identity": self.identity, + "datasets": [asdict(dataset) for dataset in self.datasets], + } + if include_hash: + payload["sha256"] = self.sha256 + return payload + + +class StatelessDrawSampler(Sampler[LogicalDraw]): + """Yield rank-disjoint logical draw ids; sampling itself stays inside the dataset.""" + + def __init__( + self, + num_local_draws: int, + rank: int = 0, + world_size: int = 1, + start_global_draw_id: int = 0, + ) -> None: + if num_local_draws < 0: + raise ValueError("num_local_draws must be non-negative") + if world_size <= 0 or rank < 0 or rank >= world_size: + raise ValueError("rank must satisfy 0 <= rank < world_size") + self.num_local_draws = int(num_local_draws) + self.rank = int(rank) + self.world_size = int(world_size) + self.start_global_draw_id = int(start_global_draw_id) + if self.start_global_draw_id < 0: + raise ValueError("start_global_draw_id must be non-negative") + self.epoch = 0 + self._resume_epoch = 0 + + def set_resume_position(self, epoch: int, start_global_draw_id: int) -> None: + """Resume at an inclusive global draw cursor within one epoch.""" + if epoch < 0 or start_global_draw_id < 0: + raise ValueError("resume epoch and draw id must be non-negative") + self.epoch = int(epoch) + self._resume_epoch = int(epoch) + self.start_global_draw_id = int(start_global_draw_id) + + def set_epoch(self, epoch: int) -> None: + if epoch < 0: + raise ValueError("epoch must be non-negative") + self.epoch = int(epoch) + if self.epoch != self._resume_epoch: + self.start_global_draw_id = 0 + + def _first_local_draw_id(self) -> int: + return max( + 0, + (self.start_global_draw_id - self.rank + self.world_size - 1) + // self.world_size, + ) + + def __iter__(self) -> Iterator[LogicalDraw]: + for local_draw_id in range(self._first_local_draw_id(), self.num_local_draws): + yield LogicalDraw( + epoch=self.epoch, + global_draw_id=local_draw_id * self.world_size + self.rank, + ) + + def __len__(self) -> int: + return max(0, self.num_local_draws - self._first_local_draw_id()) diff --git a/real_robot/source/retime/scoring/__init__.py b/real_robot/source/retime/scoring/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cfdc164b8adcc37e155501f1366842f9003d0b36 --- /dev/null +++ b/real_robot/source/retime/scoring/__init__.py @@ -0,0 +1,5 @@ +"""Offline scoring utilities for InfoFlow proposal construction.""" + +from .cnf_likelihood import CNFLikelihoodResult, conditional_cnf_log_likelihood + +__all__ = ["CNFLikelihoodResult", "conditional_cnf_log_likelihood"] diff --git a/real_robot/source/retime/scoring/cheap_proposals.py b/real_robot/source/retime/scoring/cheap_proposals.py new file mode 100644 index 0000000000000000000000000000000000000000..45e94f6549ea4652d4f0f88f68e50d63572110a6 --- /dev/null +++ b/real_robot/source/retime/scoring/cheap_proposals.py @@ -0,0 +1,210 @@ +"""Frozen low-cost proposal scorers for the prospective InfoFlow study.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from retime.scoring.residual_cache import percentile_intensity_masses + + +ACTION_CHANGE_SPEC = { + "version": 1, + "coordinate_system": "mixture_global_min_max_first_6_gripper_identity", + "score": "mean_adjacent_l2_at_anchor", + "endpoint": "single_sided", + "percentile_scope": "within_episode", +} +BRIDGERT_ACTION_CHANGE_SPEC = { + "version": 1, + "coordinate_system": "per_dataset_q01_q99_first_6_gripper_identity", + "score": "mean_adjacent_l2_at_anchor", + "endpoint": "single_sided", + "percentile_scope": "within_episode", +} +RANDOM_SPEC = { + "version": 1, + "generator": "numpy_pcg64_seeded_by_sha256", + "percentile_scope": "within_episode", +} + + +def canonical_sha256(payload: Any) -> str: + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def normalize_libero_actions( + actions: np.ndarray, + statistics: dict[str, Any], +) -> np.ndarray: + """Match the LIBERO training transform for six continuous axes + gripper.""" + values = np.asarray(actions, dtype=np.float64) + if values.ndim != 2 or values.shape[1] != 7: + raise ValueError("LIBERO actions must have shape [time, 7]") + minimum = np.asarray(statistics["min"], dtype=np.float64) + maximum = np.asarray(statistics["max"], dtype=np.float64) + if minimum.shape != (7,) or maximum.shape != (7,): + raise ValueError("LIBERO action statistics must contain seven dimensions") + if not np.isfinite(values).all() or not np.isfinite(minimum).all() or not np.isfinite(maximum).all(): + raise ValueError("actions and statistics must be finite") + + result = values.copy() + span = maximum[:6] - minimum[:6] + nonconstant = span != 0.0 + result[:, :6] = 0.0 + result[:, :6][:, nonconstant] = ( + 2.0 + * (values[:, :6][:, nonconstant] - minimum[:6][nonconstant]) + / span[nonconstant] + - 1.0 + ) + return result + + +def normalize_q99_actions( + actions: np.ndarray, + statistics: dict[str, Any], +) -> np.ndarray: + """Apply the Bridge/RT-1 training scale to continuous action axes.""" + values = np.asarray(actions, dtype=np.float64) + if values.ndim != 2 or values.shape[1] != 7: + raise ValueError("Bridge/RT-1 actions must have shape [time, 7]") + lower = np.asarray(statistics["q01"], dtype=np.float64) + upper = np.asarray(statistics["q99"], dtype=np.float64) + if lower.shape != (7,) or upper.shape != (7,): + raise ValueError("Bridge/RT-1 action statistics must contain seven dimensions") + + result = values.copy() + result[:, :6] = ( + 2.0 * (values[:, :6] - lower[:6]) / (upper[:6] - lower[:6]) - 1.0 + ) + result[:, :6] = np.clip(result[:, :6], -1.0, 1.0) + return result + + +def action_change_scores( + normalized_actions: np.ndarray, + anchors: np.ndarray, +) -> np.ndarray: + """Score local motion change around every valid training anchor.""" + actions = np.asarray(normalized_actions, dtype=np.float64) + anchor_indices = np.asarray(anchors, dtype=np.int64) + if actions.ndim != 2 or actions.shape[0] == 0: + raise ValueError("normalized_actions must be a non-empty matrix") + if anchor_indices.ndim != 1 or anchor_indices.size == 0: + raise ValueError("anchors must be a non-empty vector") + if np.any(anchor_indices < 0) or np.any(anchor_indices >= len(actions)): + raise ValueError("anchor lies outside the action trajectory") + if len(actions) == 1: + return np.zeros(len(anchor_indices), dtype=np.float64) + + deltas = np.linalg.norm(np.diff(actions, axis=0), axis=1) + scores = np.empty(len(anchor_indices), dtype=np.float64) + for output_index, anchor in enumerate(anchor_indices): + if anchor == 0: + scores[output_index] = deltas[0] + elif anchor == len(actions) - 1: + scores[output_index] = deltas[-1] + else: + scores[output_index] = 0.5 * ( + deltas[anchor - 1] + deltas[anchor] + ) + return scores + + +def random_episode_masses( + anchor_count: int, + *, + seed: int, + dataset_name: str, + trajectory_id: int | str, + epsilon: float, +) -> np.ndarray: + """Generate a fixed iid-random proposal without depending on iteration order.""" + if anchor_count <= 0: + raise ValueError("anchor_count must be positive") + identity = { + "spec": RANDOM_SPEC, + "seed": int(seed), + "dataset_name": str(dataset_name), + "trajectory_id": trajectory_id, + } + stateless_seed = int.from_bytes( + hashlib.sha256( + json.dumps(identity, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).digest()[:8], + byteorder="little", + signed=False, + ) + scores = np.random.default_rng(stateless_seed).random(anchor_count) + return percentile_intensity_masses(scores, epsilon=epsilon) + + +def load_action_statistics( + dataset_root: Path, + *, + filename: str = "stats_gr00t.json", +) -> dict[str, Any]: + stats_path = dataset_root / "meta" / filename + payload = json.loads(stats_path.read_text(encoding="utf-8")) + statistics = payload.get("statistics", payload) + if "action" not in statistics: + raise ValueError(f"Missing action statistics in {stats_path}") + return statistics["action"] + + +def merge_libero_action_statistics( + per_dataset_statistics: list[dict[str, Any]], +) -> dict[str, list[float]]: + """Match LeRobotMixtureDataset's extrema merge used by min-max training.""" + if not per_dataset_statistics: + raise ValueError("at least one dataset statistics record is required") + minima = np.asarray( + [record["min"] for record in per_dataset_statistics], + dtype=np.float64, + ) + maxima = np.asarray( + [record["max"] for record in per_dataset_statistics], + dtype=np.float64, + ) + if ( + minima.ndim != 2 + or maxima.shape != minima.shape + or minima.shape[1] != 7 + or not np.isfinite(minima).all() + or not np.isfinite(maxima).all() + ): + raise ValueError("dataset action extrema must have shape [datasets, 7]") + return { + "min": minima.min(axis=0).tolist(), + "max": maxima.max(axis=0).tolist(), + } + + +def episode_parquet_path( + dataset_root: Path, + trajectory_id: int | str, +) -> Path: + try: + episode_id = int(trajectory_id) + except (TypeError, ValueError) as error: + raise ValueError(f"Non-integer LIBERO trajectory id: {trajectory_id!r}") from error + matches = list( + (dataset_root / "data").glob( + f"chunk-*/episode_{episode_id:06d}.parquet" + ) + ) + if len(matches) != 1: + raise FileNotFoundError( + f"Expected one parquet for episode {episode_id}, found {len(matches)}" + ) + return matches[0] diff --git a/real_robot/source/retime/scoring/cnf_likelihood.py b/real_robot/source/retime/scoring/cnf_likelihood.py new file mode 100644 index 0000000000000000000000000000000000000000..175a19ff0c69399919c2a7fe11222b825eeb9a76 --- /dev/null +++ b/real_robot/source/retime/scoring/cnf_likelihood.py @@ -0,0 +1,289 @@ +"""Conditional CNF likelihood estimation for frozen action vector fields.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Callable + +import torch + + +VectorField = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + + +@dataclass(frozen=True) +class CNFLikelihoodResult: + log_probability: torch.Tensor + normalized_nll: torch.Tensor + base_latent: torch.Tensor + divergence_integral: torch.Tensor + divergence_method: str + probe_count: int + vjp_chunk_size: int + solver: str + function_evaluations: int + + +def _rademacher_probes( + shape: torch.Size, + *, + count: int, + device: torch.device, + dtype: torch.dtype, + seed: int, +) -> torch.Tensor: + generator = torch.Generator(device=device) + generator.manual_seed(int(seed)) + bits = torch.randint( + 0, + 2, + (count, 1, *shape[1:]), + generator=generator, + device=device, + dtype=torch.int8, + ) + probes = bits.to(dtype=dtype).mul_(2).sub_(1) + return probes.expand(count, int(shape[0]), *shape[1:]) + + +def _orthogonal_probes( + shape: torch.Size, + *, + count: int, + device: torch.device, + dtype: torch.dtype, + seed: int, +) -> torch.Tensor: + """Draw scaled orthonormal probes with E[mean(v v^T)] = I.""" + if len(shape) < 2: + raise ValueError("probe shape must contain batch and feature dimensions") + batch_size = int(shape[0]) + dimensions = math.prod(shape[1:]) + if count > dimensions: + raise ValueError( + f"orthogonal probe count {count} exceeds action dimension {dimensions}" + ) + + # Common random numbers make a sample's score independent of how samples + # are grouped into batches and reduce noise in pairwise ranking. + generator = torch.Generator(device="cpu") + generator.manual_seed(int(seed)) + gaussian = torch.randn( + (dimensions, count), generator=generator, dtype=torch.float64 + ) + orthonormal, _ = torch.linalg.qr(gaussian, mode="reduced") + flattened = ( + orthonormal.T[:, None, :] + .expand(count, batch_size, dimensions) + .mul(math.sqrt(float(dimensions))) + ) + return flattened.reshape(count, *shape).to(device=device, dtype=dtype) + + +def _exact_basis_probes( + shape: torch.Size, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Return a complete scaled coordinate basis whose mean gives exact trace.""" + if len(shape) < 2: + raise ValueError("probe shape must contain batch and feature dimensions") + batch_size = int(shape[0]) + dimensions = math.prod(shape[1:]) + basis = torch.eye(dimensions, device=device, dtype=dtype) + basis = basis.mul(math.sqrt(float(dimensions))) + return ( + basis[:, None, :] + .expand(dimensions, batch_size, dimensions) + .reshape(dimensions, *shape) + ) + + +def _make_probes( + shape: torch.Size, + *, + method: str, + count: int, + device: torch.device, + dtype: torch.dtype, + seed: int, +) -> torch.Tensor: + if method == "rademacher": + return _rademacher_probes( + shape, count=count, device=device, dtype=dtype, seed=seed + ) + if method == "orthogonal": + return _orthogonal_probes( + shape, count=count, device=device, dtype=dtype, seed=seed + ) + if method == "exact": + return _exact_basis_probes(shape, device=device, dtype=dtype) + raise ValueError( + "divergence_method must be one of: rademacher, orthogonal, exact" + ) + + +def _velocity_and_divergence( + vector_field: VectorField, + actions: torch.Tensor, + time: torch.Tensor, + probes: torch.Tensor, + vjp_chunk_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + with torch.enable_grad(): + differentiable_actions = actions.detach().requires_grad_(True) + velocity = vector_field(differentiable_actions, time) + if velocity.shape != differentiable_actions.shape: + raise ValueError( + "vector field output must match actions: " + f"got {tuple(velocity.shape)}, expected {tuple(actions.shape)}" + ) + + estimates = [] + for chunk_start in range(0, len(probes), vjp_chunk_size): + probe_chunk = probes[chunk_start : chunk_start + vjp_chunk_size] + chunk_end = min(chunk_start + vjp_chunk_size, len(probes)) + if len(probe_chunk) == 1: + vector_jacobian_products = torch.autograd.grad( + outputs=(velocity * probe_chunk[0]).sum(), + inputs=differentiable_actions, + retain_graph=chunk_end < len(probes), + create_graph=False, + )[0].unsqueeze(0) + else: + vector_jacobian_products = torch.autograd.grad( + outputs=velocity, + inputs=differentiable_actions, + grad_outputs=probe_chunk, + retain_graph=chunk_end < len(probes), + create_graph=False, + is_grads_batched=True, + )[0] + estimates.append( + (vector_jacobian_products * probe_chunk) + .flatten(start_dim=2) + .sum(dim=2) + ) + divergence = torch.cat(estimates, dim=0).mean(dim=0) + return velocity.detach(), divergence.detach().to(dtype=torch.float64) + + +def conditional_cnf_log_likelihood( + vector_field: VectorField, + target_actions: torch.Tensor, + *, + integration_steps: int = 16, + hutchinson_samples: int = 4, + seed: int = 0, + divergence_method: str = "rademacher", + vjp_chunk_size: int = 1, + solver: str = "rk4", +) -> CNFLikelihoodResult: + """Integrate data at ``t=1`` back to a standard-normal base at ``t=0``. + + The augmented backward ODE integrates ``dx/dt = v(x,t|c)`` and + ``dI/dt = div_x v`` from 1 to 0. Therefore + ``log p_1(x_1|c) = log p_0(x_0) + I``. + """ + + if not isinstance(target_actions, torch.Tensor): + raise TypeError("target_actions must be a torch.Tensor") + if target_actions.ndim < 2 or target_actions.shape[0] == 0: + raise ValueError("target_actions must have a non-empty batch dimension") + if not target_actions.dtype.is_floating_point: + raise TypeError("target_actions must have a floating-point dtype") + if integration_steps <= 0 or hutchinson_samples <= 0: + raise ValueError("integration_steps and hutchinson_samples must be positive") + if vjp_chunk_size <= 0: + raise ValueError("vjp_chunk_size must be positive") + if solver not in {"euler", "midpoint", "rk4"}: + raise ValueError("solver must be one of: euler, midpoint, rk4") + if not torch.isfinite(target_actions).all(): + raise ValueError("target_actions must be finite") + + actions = target_actions.detach() + batch_size = actions.shape[0] + probes = _make_probes( + actions.shape, + method=divergence_method, + count=hutchinson_samples, + device=actions.device, + dtype=actions.dtype, + seed=seed, + ) + divergence_integral = torch.zeros( + batch_size, device=actions.device, dtype=torch.float64 + ) + dt = -1.0 / float(integration_steps) + + def derivative( + state: torch.Tensor, scalar_time: float + ) -> tuple[torch.Tensor, torch.Tensor]: + time = torch.full( + (batch_size,), + scalar_time, + device=state.device, + dtype=state.dtype, + ) + return _velocity_and_divergence( + vector_field, state, time, probes, vjp_chunk_size + ) + + for step in range(integration_steps): + time = 1.0 + step * dt + if solver == "euler": + velocity, divergence = derivative(actions, time) + actions = actions + dt * velocity + divergence_integral = divergence_integral + dt * divergence + elif solver == "midpoint": + velocity_1, _ = derivative(actions, time) + velocity_2, divergence_2 = derivative( + actions + (dt / 2.0) * velocity_1, time + dt / 2.0 + ) + actions = actions + dt * velocity_2 + divergence_integral = divergence_integral + dt * divergence_2 + else: + velocity_1, divergence_1 = derivative(actions, time) + velocity_2, divergence_2 = derivative( + actions + (dt / 2.0) * velocity_1, time + dt / 2.0 + ) + velocity_3, divergence_3 = derivative( + actions + (dt / 2.0) * velocity_2, time + dt / 2.0 + ) + velocity_4, divergence_4 = derivative( + actions + dt * velocity_3, time + dt + ) + actions = actions + (dt / 6.0) * ( + velocity_1 + 2.0 * velocity_2 + 2.0 * velocity_3 + velocity_4 + ) + divergence_integral = divergence_integral + (dt / 6.0) * ( + divergence_1 + + 2.0 * divergence_2 + + 2.0 * divergence_3 + + divergence_4 + ) + + dimensions = math.prod(target_actions.shape[1:]) + base_latent64 = actions.to(dtype=torch.float64) + base_log_probability = -0.5 * ( + base_latent64.square() + math.log(2.0 * math.pi) + ).flatten(start_dim=1).sum(dim=1) + log_probability = base_log_probability + divergence_integral + normalized_nll = -log_probability / float(dimensions) + if not torch.isfinite(log_probability).all(): + raise FloatingPointError("CNF likelihood integration produced non-finite values") + + return CNFLikelihoodResult( + log_probability=log_probability, + normalized_nll=normalized_nll, + base_latent=actions, + divergence_integral=divergence_integral, + divergence_method=divergence_method, + probe_count=int(probes.shape[0]), + vjp_chunk_size=min(int(vjp_chunk_size), int(probes.shape[0])), + solver=solver, + function_evaluations=integration_steps + * {"euler": 1, "midpoint": 2, "rk4": 4}[solver], + ) diff --git a/real_robot/source/retime/scoring/discrete_flow_likelihood.py b/real_robot/source/retime/scoring/discrete_flow_likelihood.py new file mode 100644 index 0000000000000000000000000000000000000000..0b7d64e19e72820db21e1d23b1164af8aee82188 --- /dev/null +++ b/real_robot/source/retime/scoring/discrete_flow_likelihood.py @@ -0,0 +1,185 @@ +"""Likelihood of the finite-step Euler map used by the StarVLA action head.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Callable + +import torch + + +VectorField = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + + +@dataclass(frozen=True) +class DiscreteFlowLikelihoodResult: + log_probability: torch.Tensor + normalized_nll: torch.Tensor + base_latent: torch.Tensor + logabsdet_sum: torch.Tensor + inverse_iterations: tuple[int, ...] + inverse_residual_max: torch.Tensor + minimum_singular_value: torch.Tensor + + +def _velocity_jacobian( + vector_field: VectorField, + actions: torch.Tensor, + time: torch.Tensor, + *, + chunk_size: int, +) -> torch.Tensor: + """Materialize the per-example Jacobian with chunked batched VJPs.""" + batch_size = actions.shape[0] + dimensions = math.prod(actions.shape[1:]) + differentiable_actions = actions.detach().requires_grad_(True) + with torch.enable_grad(): + velocity = vector_field(differentiable_actions, time) + if velocity.shape != differentiable_actions.shape: + raise ValueError("vector field output must match action shape") + basis = torch.eye(dimensions, device=actions.device, dtype=actions.dtype) + rows = [] + for start in range(0, dimensions, chunk_size): + basis_chunk = basis[start : start + chunk_size] + grad_outputs = ( + basis_chunk[:, None, :] + .expand(len(basis_chunk), batch_size, dimensions) + .reshape(len(basis_chunk), *actions.shape) + ) + rows.append( + torch.autograd.grad( + outputs=velocity, + inputs=differentiable_actions, + grad_outputs=grad_outputs, + retain_graph=start + chunk_size < dimensions, + create_graph=False, + is_grads_batched=True, + )[0].flatten(start_dim=2) + ) + return torch.cat(rows, dim=0).permute(1, 0, 2).detach() + + +def _inverse_euler_step( + vector_field: VectorField, + target: torch.Tensor, + time: torch.Tensor, + *, + dt: float, + tolerance: float, + max_iterations: int, +) -> tuple[torch.Tensor, int, torch.Tensor]: + """Solve x + dt*v(x,t) = target by fixed-point iteration.""" + estimate = target.detach() + converged_at = None + with torch.no_grad(): + for iteration in range(1, max_iterations + 1): + updated = target - dt * vector_field(estimate, time) + delta = (updated - estimate).flatten(start_dim=1).abs().amax(dim=1) + estimate = updated + if bool(torch.all(delta <= tolerance)): + converged_at = iteration + break + forward_residual = ( + estimate + dt * vector_field(estimate, time) - target + ).flatten(start_dim=1).abs().amax(dim=1) + if converged_at is None or not bool(torch.all(forward_residual <= tolerance)): + raise RuntimeError( + "Euler residual map inversion did not converge: " + f"max residual={forward_residual.max().item():.6g}" + ) + return estimate, converged_at, forward_residual + + +def conditional_discrete_flow_log_likelihood( + vector_field: VectorField, + target_actions: torch.Tensor, + *, + num_steps: int = 4, + inverse_tolerance: float = 1e-5, + max_inverse_iterations: int = 50, + jacobian_chunk_size: int = 16, +) -> DiscreteFlowLikelihoodResult: + """Invert the deployed Euler sampler and apply exact change of variables.""" + if not isinstance(target_actions, torch.Tensor): + raise TypeError("target_actions must be a torch.Tensor") + if target_actions.ndim < 2 or target_actions.shape[0] == 0: + raise ValueError("target_actions must have a non-empty batch dimension") + if not target_actions.dtype.is_floating_point: + raise TypeError("target_actions must have a floating-point dtype") + if not torch.isfinite(target_actions).all(): + raise ValueError("target_actions must be finite") + if num_steps <= 0 or max_inverse_iterations <= 0 or jacobian_chunk_size <= 0: + raise ValueError("step, iteration, and chunk counts must be positive") + if inverse_tolerance <= 0: + raise ValueError("inverse_tolerance must be positive") + + actions = target_actions.detach() + batch_size = actions.shape[0] + dimensions = math.prod(actions.shape[1:]) + dt = 1.0 / float(num_steps) + logabsdet_sum = torch.zeros( + batch_size, device=actions.device, dtype=torch.float64 + ) + minimum_singular_value = torch.full( + (batch_size,), float("inf"), device=actions.device, dtype=torch.float64 + ) + residual_max = torch.zeros( + batch_size, device=actions.device, dtype=torch.float64 + ) + inverse_iterations = [] + identity = torch.eye(dimensions, device=actions.device, dtype=actions.dtype) + + for step in reversed(range(num_steps)): + time = torch.full( + (batch_size,), + step / float(num_steps), + device=actions.device, + dtype=actions.dtype, + ) + actions, iterations, residual = _inverse_euler_step( + vector_field, + actions, + time, + dt=dt, + tolerance=inverse_tolerance, + max_iterations=max_inverse_iterations, + ) + inverse_iterations.append(iterations) + residual_max = torch.maximum(residual_max, residual.to(torch.float64)) + + velocity_jacobian = _velocity_jacobian( + vector_field, + actions, + time, + chunk_size=jacobian_chunk_size, + ) + step_jacobian = identity[None] + dt * velocity_jacobian + sign, logabsdet = torch.linalg.slogdet(step_jacobian) + if bool(torch.any(sign == 0)) or not bool(torch.all(torch.isfinite(logabsdet))): + raise RuntimeError("Euler residual map has a singular Jacobian") + singular_values = torch.linalg.svdvals(step_jacobian) + minimum_singular_value = torch.minimum( + minimum_singular_value, + singular_values[:, -1].to(torch.float64), + ) + logabsdet_sum = logabsdet_sum + logabsdet.to(torch.float64) + + base_latent64 = actions.to(torch.float64) + base_log_probability = -0.5 * ( + base_latent64.square() + math.log(2.0 * math.pi) + ).flatten(start_dim=1).sum(dim=1) + log_probability = base_log_probability - logabsdet_sum + normalized_nll = -log_probability / float(dimensions) + if not bool(torch.all(torch.isfinite(log_probability))): + raise FloatingPointError("Discrete-flow likelihood produced non-finite values") + + return DiscreteFlowLikelihoodResult( + log_probability=log_probability, + normalized_nll=normalized_nll, + base_latent=actions, + logabsdet_sum=logabsdet_sum, + inverse_iterations=tuple(inverse_iterations), + inverse_residual_max=residual_max, + minimum_singular_value=minimum_singular_value, + ) diff --git a/real_robot/source/retime/scoring/flow_residual.py b/real_robot/source/retime/scoring/flow_residual.py new file mode 100644 index 0000000000000000000000000000000000000000..3a1a8093a4a340c7bc61e144a876b04854f80c2a --- /dev/null +++ b/real_robot/source/retime/scoring/flow_residual.py @@ -0,0 +1,265 @@ +"""Low-variance frozen Flow-Matching residual estimators.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, Sequence + +import numpy as np +import torch +from scipy.stats import beta as scipy_beta + + +@dataclass(frozen=True) +class FlowResidualResult: + """Residual estimates for one conditioned action batch.""" + + mean_residual: torch.Tensor + estimator_residuals: torch.Tensor + sample_losses: torch.Tensor + time_values: torch.Tensor + estimator_seeds: tuple[int, ...] + samples_per_estimator: int + + +@lru_cache(maxsize=32) +def _deterministic_paired_times( + sample_count: int, + beta_alpha: float, + beta_beta: float, + noise_s: float, +) -> tuple[float, ...]: + pair_count = sample_count // 2 + quantiles = (np.arange(pair_count, dtype=np.float64) + 0.5) / pair_count + beta_samples = scipy_beta.ppf(quantiles, a=beta_alpha, b=beta_beta) + beta_samples = np.minimum(beta_samples, noise_s) + times = (noise_s - beta_samples) / noise_s + return tuple(float(value) for value in np.repeat(times, 2)) + + +def antithetic_stratified_flow_samples( + actions: torch.Tensor, + *, + sample_count: int, + seed: int, + beta_alpha: float, + beta_beta: float, + noise_s: float, + shared_noise_across_batch: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate paired Gaussian noise at deterministic Beta time quantiles.""" + if actions.ndim < 2 or not torch.is_floating_point(actions): + raise ValueError("actions must be a floating-point tensor with a batch axis") + if sample_count < 2 or sample_count % 2: + raise ValueError("sample_count must be a positive even integer") + if beta_alpha <= 0.0 or beta_beta <= 0.0: + raise ValueError("Beta shape parameters must be positive") + if not 0.0 < noise_s <= 1.0: + raise ValueError("noise_s must satisfy 0 < noise_s <= 1") + + pair_count = sample_count // 2 + paired_times = _deterministic_paired_times( + sample_count, + float(beta_alpha), + float(beta_beta), + float(noise_s), + ) + + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + noise_batch_size = 1 if shared_noise_across_batch else actions.shape[0] + base_noise = torch.randn( + (pair_count, noise_batch_size, *actions.shape[1:]), + generator=generator, + dtype=torch.float32, + device="cpu", + ) + paired_noise = torch.stack((base_noise, -base_noise), dim=1).reshape( + sample_count, + noise_batch_size, + *actions.shape[1:], + ) + if shared_noise_across_batch: + paired_noise = paired_noise.expand(sample_count, *actions.shape) + return ( + paired_noise.to(device=actions.device, dtype=actions.dtype), + torch.as_tensor( + paired_times, + device=actions.device, + dtype=actions.dtype, + ), + ) + + +def randomized_antithetic_stratified_flow_samples( + actions: torch.Tensor, + *, + sample_count: int, + seed: int, + beta_alpha: float, + beta_beta: float, + noise_s: float, + shared_noise_across_batch: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate an unbiased randomized-stratified time/noise draw. + + One uniform quantile is drawn inside each equal-probability Beta stratum. + Gaussian noises are paired as ``z`` and ``-z`` at the same sampled time. + Averaging all returned samples is unbiased for the joint time/noise + expectation while retaining variance reduction from stratification and + antithetic pairing. + """ + if actions.ndim < 2 or not torch.is_floating_point(actions): + raise ValueError("actions must be a floating-point tensor with a batch axis") + if sample_count < 2 or sample_count % 2: + raise ValueError("sample_count must be a positive even integer") + if beta_alpha <= 0.0 or beta_beta <= 0.0: + raise ValueError("Beta shape parameters must be positive") + if not 0.0 < noise_s <= 1.0: + raise ValueError("noise_s must satisfy 0 < noise_s <= 1") + + pair_count = sample_count // 2 + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + jitter = torch.rand( + pair_count, + generator=generator, + dtype=torch.float64, + device="cpu", + ).numpy() + quantiles = (np.arange(pair_count, dtype=np.float64) + jitter) / pair_count + beta_samples = scipy_beta.ppf( + quantiles, + a=beta_alpha, + b=beta_beta, + ) + beta_samples = np.minimum(beta_samples, noise_s) + times = (noise_s - beta_samples) / noise_s + paired_times = np.repeat(times, 2) + + noise_batch_size = 1 if shared_noise_across_batch else actions.shape[0] + base_noise = torch.randn( + (pair_count, noise_batch_size, *actions.shape[1:]), + generator=generator, + dtype=torch.float32, + device="cpu", + ) + paired_noise = torch.stack((base_noise, -base_noise), dim=1).reshape( + sample_count, + noise_batch_size, + *actions.shape[1:], + ) + if shared_noise_across_batch: + paired_noise = paired_noise.expand(sample_count, *actions.shape) + return ( + paired_noise.to(device=actions.device, dtype=actions.dtype), + torch.as_tensor( + paired_times, + device=actions.device, + dtype=actions.dtype, + ), + ) + + +def flow_matching_residual_samples( + vector_field: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + actions: torch.Tensor, + noises: torch.Tensor, + times: torch.Tensor, +) -> torch.Tensor: + """Return per-Monte-Carlo, per-example mean-squared FM residuals.""" + if actions.ndim != 3: + raise ValueError("actions must have shape [batch, horizon, action_dim]") + if noises.ndim != 4 or noises.shape[1:] != actions.shape: + raise ValueError( + "noises must have shape [samples, batch, horizon, action_dim]" + ) + if times.shape != (noises.shape[0],): + raise ValueError("times must have shape [samples]") + + sample_count, batch_size, horizon, action_dim = noises.shape + expanded_actions = actions.unsqueeze(0).expand_as(noises) + expanded_times = times[:, None, None, None] + noisy_trajectories = ( + (1.0 - expanded_times) * noises + expanded_times * expanded_actions + ) + target_velocity = expanded_actions - noises + flat_noisy_trajectories = noisy_trajectories.reshape( + sample_count * batch_size, + horizon, + action_dim, + ) + flat_times = times[:, None].expand(sample_count, batch_size).reshape(-1) + predicted_velocity = vector_field(flat_noisy_trajectories, flat_times) + if predicted_velocity.shape != flat_noisy_trajectories.shape: + raise ValueError("vector_field returned an unexpected shape") + # The vector-field output is private to this scoring call. Reusing its + # storage avoids another sample_count-sized tensor at the DiT peak. + elementwise_loss = predicted_velocity.reshape_as(noisy_trajectories) + elementwise_loss.sub_(target_velocity).square_() + return elementwise_loss.mean(dim=(-1, -2)) + + +def estimate_flow_matching_residual( + vector_field: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + actions: torch.Tensor, + *, + estimator_seeds: Sequence[int], + samples_per_estimator: int, + beta_alpha: float, + beta_beta: float, + noise_s: float, + shared_noise_across_batch: bool = False, +) -> FlowResidualResult: + """Estimate residual with independent antithetic, time-stratified replicates.""" + seeds = tuple(int(seed) for seed in estimator_seeds) + if len(seeds) < 2 or len(set(seeds)) != len(seeds): + raise ValueError("estimator_seeds must contain at least two unique seeds") + + # Preserve the historical CPU RNG protocol exactly, but concatenate on the + # CPU and issue one host-to-device copy instead of two copies per seed. + sample_actions = ( + torch.empty(actions.shape, dtype=actions.dtype, device="cpu") + if actions.device.type != "cpu" + else actions + ) + noises = [] + times = [] + for seed in seeds: + seed_noises, seed_times = antithetic_stratified_flow_samples( + sample_actions, + sample_count=samples_per_estimator, + seed=seed, + beta_alpha=beta_alpha, + beta_beta=beta_beta, + noise_s=noise_s, + shared_noise_across_batch=shared_noise_across_batch, + ) + noises.append(seed_noises) + times.append(seed_times) + + fused_noises = torch.cat(noises, dim=0).to(device=actions.device) + fused_times = torch.cat(times, dim=0).to(device=actions.device) + + fused_losses = flow_matching_residual_samples( + vector_field, + actions, + fused_noises, + fused_times, + ) + sample_losses = fused_losses.reshape( + len(seeds), + samples_per_estimator, + actions.shape[0], + ) + estimator_residuals = sample_losses.mean(dim=1) + time_values = fused_times[:samples_per_estimator] + return FlowResidualResult( + mean_residual=estimator_residuals.mean(dim=0), + estimator_residuals=estimator_residuals, + sample_losses=sample_losses, + time_values=time_values, + estimator_seeds=seeds, + samples_per_estimator=samples_per_estimator, + ) diff --git a/real_robot/source/retime/scoring/partial_gradient.py b/real_robot/source/retime/scoring/partial_gradient.py new file mode 100644 index 0000000000000000000000000000000000000000..4310619529ec7afa270cec20beb143397ee97b96 --- /dev/null +++ b/real_robot/source/retime/scoring/partial_gradient.py @@ -0,0 +1,114 @@ +"""Parameter-group helpers for partial-gradient importance proxies.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn + + +@dataclass(frozen=True) +class PartialGradientGroup: + """A named, non-empty subset of action-model parameters.""" + + name: str + parameter_names: tuple[str, ...] + parameter_count: int + + +def define_partial_gradient_groups( + action_model: nn.Module, +) -> dict[str, PartialGradientGroup]: + """Return the four parameter groups used for partial-gradient analysis.""" + named_parameters = dict(action_model.named_parameters()) + if not named_parameters: + raise ValueError("action model has no parameters") + + transformer_blocks = getattr( + getattr(action_model, "model", None), + "transformer_blocks", + None, + ) + if transformer_blocks is None or len(transformer_blocks) == 0: + raise ValueError("action model has no DiT transformer blocks") + last_block_prefix = f"model.transformer_blocks.{len(transformer_blocks) - 1}." + + selectors = { + "decoder": lambda name: name.startswith("action_decoder."), + "output_head": lambda name: ( + name.startswith("action_decoder.") + or name.startswith("model.proj_out_1.") + or name.startswith("model.proj_out_2.") + ), + "last_block_and_head": lambda name: ( + name.startswith(last_block_prefix) + or name.startswith("action_decoder.") + or name.startswith("model.proj_out_1.") + or name.startswith("model.proj_out_2.") + ), + "normalization": lambda name: "norm" in name.lower(), + } + + groups = {} + for group_name, selector in selectors.items(): + selected_names = tuple( + name for name in named_parameters if selector(name) + ) + if not selected_names: + raise ValueError( + f"partial-gradient group {group_name!r} has no parameters" + ) + groups[group_name] = PartialGradientGroup( + name=group_name, + parameter_names=selected_names, + parameter_count=sum( + named_parameters[name].numel() for name in selected_names + ), + ) + return groups + + +def activate_partial_gradient_group( + action_model: nn.Module, + group: PartialGradientGroup, +) -> list[nn.Parameter]: + """Freeze the action model except for one exact, validated group.""" + named_parameters = dict(action_model.named_parameters()) + requested = set(group.parameter_names) + missing = requested - set(named_parameters) + if missing: + raise ValueError( + f"group {group.name!r} contains unknown parameters: {sorted(missing)}" + ) + + action_model.requires_grad_(False) + selected = [] + for name, parameter in named_parameters.items(): + if name in requested: + parameter.requires_grad_(True) + selected.append(parameter) + if not selected: + raise RuntimeError(f"group {group.name!r} activated no parameters") + observed_count = sum(parameter.numel() for parameter in selected) + if observed_count != group.parameter_count: + raise RuntimeError( + f"group {group.name!r} parameter count changed from " + f"{group.parameter_count} to {observed_count}" + ) + return selected + + +def finite_gradient_l2_norm(parameters: list[nn.Parameter]) -> float: + """Compute an exact finite L2 norm over the gradients that were produced.""" + component_norms = [] + for parameter in parameters: + if parameter.grad is None: + continue + gradient = parameter.grad.detach() + if not torch.isfinite(gradient).all(): + raise FloatingPointError("partial-gradient proxy is non-finite") + component_norms.append(torch.linalg.vector_norm(gradient.float())) + if not component_norms: + raise RuntimeError("partial-gradient group received no gradients") + return float(torch.linalg.vector_norm(torch.stack(component_norms)).item()) diff --git a/real_robot/source/retime/scoring/residual_cache.py b/real_robot/source/retime/scoring/residual_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..66328863a9a64813b34119708e820b2cb2547d49 --- /dev/null +++ b/real_robot/source/retime/scoring/residual_cache.py @@ -0,0 +1,79 @@ +"""Deterministic grid scoring and interpolation for residual proposals.""" + +from __future__ import annotations + +import numpy as np +from scipy.stats import rankdata + + +def quantile_anchor_positions(anchor_count: int, grid_size: int) -> np.ndarray: + """Select fixed, unique, endpoint-covering positions within one episode.""" + if anchor_count <= 0: + raise ValueError("anchor_count must be positive") + if grid_size <= 0: + raise ValueError("grid_size must be positive") + selected_count = min(int(anchor_count), int(grid_size)) + positions = np.linspace( + 0, + anchor_count - 1, + selected_count, + dtype=np.int64, + ) + if len(np.unique(positions)) != selected_count: + raise RuntimeError("quantile anchor selection produced duplicates") + return positions + + +def interpolate_grid_scores( + anchor_count: int, + grid_positions: np.ndarray, + grid_scores: np.ndarray, +) -> np.ndarray: + """Linearly interpolate frozen grid scores over all episode anchors.""" + if anchor_count <= 0: + raise ValueError("anchor_count must be positive") + positions = np.asarray(grid_positions, dtype=np.int64) + scores = np.asarray(grid_scores, dtype=np.float64) + if ( + positions.ndim != 1 + or scores.ndim != 1 + or positions.shape != scores.shape + or positions.size == 0 + ): + raise ValueError("grid positions and scores must be aligned non-empty vectors") + if not np.isfinite(scores).all(): + raise ValueError("grid scores must be finite") + if np.any(positions < 0) or np.any(positions >= anchor_count): + raise ValueError("grid positions are outside the episode") + if np.any(np.diff(positions) <= 0): + raise ValueError("grid positions must be strictly increasing") + if positions[0] != 0 or positions[-1] != anchor_count - 1: + raise ValueError("grid positions must cover both episode endpoints") + if anchor_count == 1: + return scores.copy() + return np.interp( + np.arange(anchor_count, dtype=np.float64), + positions.astype(np.float64), + scores, + ) + + +def percentile_intensity_masses( + scores: np.ndarray, + *, + epsilon: float, +) -> np.ndarray: + """Convert episode-local scores into epsilon-floored percentile masses.""" + values = np.asarray(scores, dtype=np.float64) + if values.ndim != 1 or values.size == 0 or not np.isfinite(values).all(): + raise ValueError("scores must be a finite non-empty vector") + if epsilon <= 0.0: + raise ValueError("epsilon must be positive") + if values.size == 1: + percentiles = np.zeros(1, dtype=np.float64) + else: + percentiles = (rankdata(values, method="average") - 1.0) / ( + values.size - 1.0 + ) + return float(epsilon) + percentiles + diff --git a/real_robot/source/retime/scoring/shuffle.py b/real_robot/source/retime/scoring/shuffle.py new file mode 100644 index 0000000000000000000000000000000000000000..f95a68205a19b521f8de9ed250ee314877291bec --- /dev/null +++ b/real_robot/source/retime/scoring/shuffle.py @@ -0,0 +1,123 @@ +"""Shuffle controls for residual temporal proposals.""" + +from __future__ import annotations + +import hashlib +import json +from typing import Any + +import numpy as np + +from retime.sampler import ValidAnchorManifest +from retime.proposal import TemporalProposal + + +def _episode_seed(seed: int, dataset_index: int, episode_index: int) -> int: + payload = f"{int(seed)}:{dataset_index}:{episode_index}".encode("utf-8") + return int.from_bytes( + hashlib.sha256(payload).digest()[:8], + "little", + signed=False, + ) + + +def build_shuffle_proposal( + *, + manifest: ValidAnchorManifest, + residual_proposal: TemporalProposal, + permutation_seed: int, +) -> TemporalProposal: + """Permute residual masses within every trajectory and preserve all marginals.""" + + if residual_proposal.manifest_support != manifest.to_dict(include_hash=False)["datasets"]: + raise ValueError("residual proposal changes dataset/episode/anchor support") + if residual_proposal.score_kind != "residual": + raise ValueError("Shuffle control requires a residual source proposal") + permuted_masses = [] + for dataset_index, dataset_masses in enumerate( + residual_proposal.proposal_masses + ): + permuted_dataset = [] + for episode_index, source_mass in enumerate(dataset_masses): + source = np.asarray(source_mass, dtype=np.float64) + rng = np.random.default_rng( + _episode_seed( + permutation_seed, + dataset_index, + episode_index, + ) + ) + permuted_dataset.append(source[rng.permutation(len(source))]) + permuted_masses.append(permuted_dataset) + + identity: dict[str, Any] = { + "schema_version": 1, + "source_proposal_sha256": residual_proposal.sha256, + "source_scorer_sha256": residual_proposal.scorer_sha256, + "permutation_seed": int(permutation_seed), + "permutation_scope": "within_trajectory", + "mass_multiset_preserved": True, + "probability_multiset_preserved": True, + "importance_weight_multiset_preserved": True, + } + scorer_sha256 = hashlib.sha256( + json.dumps( + identity, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + proposal = TemporalProposal.from_masses( + manifest=manifest, + alpha=residual_proposal.alpha, + masses=permuted_masses, + score_kind="shuffle", + scorer_sha256=scorer_sha256, + validation_certificate=identity, + ) + + for dataset_index, (source_dataset, shuffled_dataset) in enumerate( + zip( + residual_proposal.proposal_masses, + proposal.proposal_masses, + ) + ): + for episode_index, (source_mass, shuffled_mass) in enumerate( + zip(source_dataset, shuffled_dataset) + ): + if not np.array_equal( + np.sort(np.asarray(source_mass)), + np.sort(np.asarray(shuffled_mass)), + ): + raise RuntimeError( + "Shuffle mass multiset differs at " + f"dataset={dataset_index} episode={episode_index}" + ) + source_q = residual_proposal.anchor_probabilities[ + dataset_index + ][episode_index] + shuffled_q = proposal.anchor_probabilities[dataset_index][ + episode_index + ] + source_w = residual_proposal.importance_weights[dataset_index][ + episode_index + ] + shuffled_w = proposal.importance_weights[dataset_index][ + episode_index + ] + if not np.allclose( + np.sort(source_q), + np.sort(shuffled_q), + rtol=0.0, + atol=1e-15, + ): + raise RuntimeError("Shuffle probability multiset differs") + if not np.allclose( + np.sort(source_w), + np.sort(shuffled_w), + rtol=0.0, + atol=1e-15, + ): + raise RuntimeError("Shuffle weight multiset differs") + + return proposal diff --git a/real_robot/source/retime/scoring/starvla_likelihood.py b/real_robot/source/retime/scoring/starvla_likelihood.py new file mode 100644 index 0000000000000000000000000000000000000000..097f7bb9f5459735b69d5c6af119a4087cfd6062 --- /dev/null +++ b/real_robot/source/retime/scoring/starvla_likelihood.py @@ -0,0 +1,165 @@ +"""Bridge StarVLA examples to the model-independent CNF likelihood integrator.""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np +import torch + +from retime.scoring.cnf_likelihood import ( + CNFLikelihoodResult, + conditional_cnf_log_likelihood, +) +from retime.scoring.discrete_flow_likelihood import ( + DiscreteFlowLikelihoodResult, + conditional_discrete_flow_log_likelihood, +) +from retime.scoring.flow_residual import ( + FlowResidualResult, + estimate_flow_matching_residual, +) + + +def _tile_conditioning(tensor: torch.Tensor | None, repeat_factor: int): + if tensor is None or repeat_factor == 1: + return tensor + tiled = tensor.unsqueeze(0).expand(repeat_factor, *tensor.shape) + return tiled.reshape(repeat_factor * tensor.shape[0], *tensor.shape[1:]) + + +def _prepare_starvla_conditioning(model, examples: Sequence[dict]): + if model.training: + raise ValueError("conditional likelihood requires model.eval() to disable dropout") + if any(parameter.requires_grad for parameter in model.parameters()): + raise ValueError("conditional likelihood requires a fully frozen model") + if not examples: + raise ValueError("examples must be non-empty") + + batch_images = [example["image"] for example in examples] + instructions = [example["lang"] for example in examples] + raw_states = ( + [example["state"] for example in examples] + if "state" in examples[0] + else None + ) + raw_actions = [example["action"] for example in examples] + + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + vlm_inputs = model.qwen_vl_interface.build_qwenvl_inputs( + images=batch_images, instructions=instructions + ) + vlm_outputs = model.qwen_vl_interface( + **vlm_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + # Both the Qwen interface and Florence's multiview adapter expose the + # action conditioning through hidden_states[-1], as in policy.forward. + hidden_states = vlm_outputs.hidden_states[-1].detach() + + actions = torch.as_tensor( + np.asarray(raw_actions), + device=hidden_states.device, + dtype=hidden_states.dtype, + )[:, -model.action_horizon :, :] + states = ( + torch.as_tensor( + np.asarray(raw_states), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + if raw_states is not None + else None + ) + + def vector_field(action_state: torch.Tensor, time: torch.Tensor) -> torch.Tensor: + requested_batch = action_state.shape[0] + conditioning_batch = hidden_states.shape[0] + if requested_batch % conditioning_batch != 0: + raise ValueError( + "action_state batch must be an integer multiple of the " + "conditioning batch" + ) + repeat_factor = requested_batch // conditioning_batch + return model.action_model.predict_velocity( + _tile_conditioning(hidden_states, repeat_factor), + action_state, + time, + state=_tile_conditioning(states, repeat_factor), + ) + return actions, vector_field + + +def score_starvla_conditional_actions( + model, + examples: Sequence[dict], + *, + integration_steps: int = 16, + hutchinson_samples: int = 4, + seed: int = 0, + divergence_method: str = "rademacher", + vjp_chunk_size: int = 1, + solver: str = "rk4", +) -> CNFLikelihoodResult: + """Score normalized expert action chunks under the continuous StarVLA ODE.""" + actions, vector_field = _prepare_starvla_conditioning(model, examples) + + return conditional_cnf_log_likelihood( + vector_field, + actions, + integration_steps=integration_steps, + hutchinson_samples=hutchinson_samples, + seed=seed, + divergence_method=divergence_method, + vjp_chunk_size=vjp_chunk_size, + solver=solver, + ) + + +def score_starvla_discrete_actions( + model, + examples: Sequence[dict], + *, + num_steps: int | None = None, + inverse_tolerance: float = 1e-5, + max_inverse_iterations: int = 50, + jacobian_chunk_size: int = 16, +) -> DiscreteFlowLikelihoodResult: + """Score actions under the finite-step Euler map used at deployment.""" + actions, vector_field = _prepare_starvla_conditioning(model, examples) + return conditional_discrete_flow_log_likelihood( + vector_field, + actions, + num_steps=( + int(model.action_model.num_inference_timesteps) + if num_steps is None + else int(num_steps) + ), + inverse_tolerance=inverse_tolerance, + max_inverse_iterations=max_inverse_iterations, + jacobian_chunk_size=jacobian_chunk_size, + ) + + +def score_starvla_flow_matching_residual( + model, + examples: Sequence[dict], + *, + estimator_seeds: Sequence[int], + samples_per_estimator: int = 8, +) -> FlowResidualResult: + """Score frozen examples with stratified-time, antithetic-noise residuals.""" + actions, vector_field = _prepare_starvla_conditioning(model, examples) + action_config = model.action_model.config + return estimate_flow_matching_residual( + vector_field, + actions, + estimator_seeds=estimator_seeds, + samples_per_estimator=samples_per_estimator, + beta_alpha=float(action_config.noise_beta_alpha), + beta_beta=float(action_config.noise_beta_beta), + noise_s=float(action_config.noise_s), + shared_noise_across_batch=True, + ) diff --git a/real_robot/source/retime/scoring/utility_audit.py b/real_robot/source/retime/scoring/utility_audit.py new file mode 100644 index 0000000000000000000000000000000000000000..9aa4cfe841f1a68e6971475227c6e8359a13d368 --- /dev/null +++ b/real_robot/source/retime/scoring/utility_audit.py @@ -0,0 +1,471 @@ +"""Pure statistics for auditing sample-importance proxy signals.""" + +from __future__ import annotations + +import numpy as np +from scipy.stats import kendalltau, rankdata, spearmanr + + +def proposal_from_scores( + scores: np.ndarray, + *, + alpha: float, + percentile_epsilon: float, +) -> tuple[np.ndarray, np.ndarray]: + """Build the empirical safe-mixture proposal used by K3 audits.""" + scores = np.asarray(scores, dtype=np.float64) + if scores.ndim != 1 or len(scores) < 2 or not np.isfinite(scores).all(): + raise ValueError("scores must be a finite one-dimensional array of length >= 2") + if not 0.0 <= alpha < 1.0: + raise ValueError("alpha must satisfy 0 <= alpha < 1") + if percentile_epsilon <= 0.0: + raise ValueError("percentile_epsilon must be positive") + + sample_count = len(scores) + baseline = np.full(sample_count, 1.0 / sample_count, dtype=np.float64) + percentiles = ( + rankdata(scores, method="average") - 1.0 + ) / (sample_count - 1.0) + intensity = percentile_epsilon + percentiles + information = baseline * intensity + information /= information.sum() + proposal = (1.0 - alpha) * baseline + alpha * information + return proposal, percentiles + + +def gradient_second_moment( + proposal: np.ndarray, + gradient_rms: np.ndarray, +) -> float: + """Return sum_i p_i^2 h_i^2 / q_i for the empirical uniform baseline.""" + proposal = np.asarray(proposal, dtype=np.float64) + gradient_rms = np.asarray(gradient_rms, dtype=np.float64) + if ( + proposal.ndim != 1 + or gradient_rms.ndim != 1 + or proposal.shape != gradient_rms.shape + or len(proposal) < 2 + ): + raise ValueError("proposal and gradient_rms must be aligned 1-D arrays") + if ( + not np.isfinite(proposal).all() + or not np.isfinite(gradient_rms).all() + or np.any(proposal <= 0.0) + or np.any(gradient_rms < 0.0) + ): + raise ValueError("proposal must be positive and gradient_rms non-negative") + if not np.isclose(proposal.sum(), 1.0, rtol=0.0, atol=1e-10): + raise ValueError("proposal must sum to one") + + baseline = np.full(len(proposal), 1.0 / len(proposal), dtype=np.float64) + return float(np.sum(np.square(baseline * gradient_rms) / proposal)) + + +def alignment_audit( + scores: np.ndarray, + gradient_rms: np.ndarray, + mean_losses: np.ndarray, + *, + alpha: float, + percentile_epsilon: float, + permutation_count: int, + permutation_seed: int, + score_label: str = "score", +) -> dict: + """Compare a proxy with gradient RMS and score-permuted proposals.""" + scores = np.asarray(scores, dtype=np.float64) + gradient_rms = np.asarray(gradient_rms, dtype=np.float64) + mean_losses = np.asarray(mean_losses, dtype=np.float64) + if ( + scores.ndim != 1 + or gradient_rms.ndim != 1 + or mean_losses.ndim != 1 + or not (scores.shape == gradient_rms.shape == mean_losses.shape) + or len(scores) < 2 + ): + raise ValueError("audit inputs must be aligned one-dimensional arrays") + if not ( + np.isfinite(scores).all() + and np.isfinite(gradient_rms).all() + and np.isfinite(mean_losses).all() + ): + raise ValueError("alignment inputs must be finite") + if np.any(gradient_rms <= 0.0): + raise ValueError("gradient_rms must be strictly positive") + if permutation_count <= 0: + raise ValueError("permutation_count must be positive") + if not score_label or not score_label.replace("_", "").isalnum(): + raise ValueError("score_label must be a non-empty identifier") + + proposal, percentiles = proposal_from_scores( + scores, + alpha=alpha, + percentile_epsilon=percentile_epsilon, + ) + baseline = np.full(len(scores), 1.0 / len(scores), dtype=np.float64) + baseline_second_moment = float(np.mean(np.square(gradient_rms))) + information_second_moment = gradient_second_moment(proposal, gradient_rms) + optimal_proposal = gradient_rms / gradient_rms.sum() + optimal_second_moment = gradient_second_moment( + optimal_proposal, + gradient_rms, + ) + + score_gradient = spearmanr(scores, gradient_rms) + score_loss = spearmanr(scores, mean_losses) + score_gradient_kendall = kendalltau(scores, gradient_rms) + if not np.isfinite(score_gradient.statistic): + raise ValueError("scores or gradient_rms have no rank variation") + top_count = max(1, len(scores) // 4) + score_top = set(np.argsort(scores)[-top_count:]) + gradient_top = set(np.argsort(gradient_rms)[-top_count:]) + + rng = np.random.default_rng(permutation_seed) + random_rhos = np.empty(permutation_count, dtype=np.float64) + random_second_moment_ratios = np.empty(permutation_count, dtype=np.float64) + for index in range(permutation_count): + shuffled_percentiles = rng.permutation(percentiles) + random_rhos[index] = spearmanr( + shuffled_percentiles, + gradient_rms, + ).statistic + random_information = baseline * ( + percentile_epsilon + shuffled_percentiles + ) + random_information /= random_information.sum() + random_proposal = ( + (1.0 - alpha) * baseline + alpha * random_information + ) + random_second_moment_ratios[index] = ( + gradient_second_moment(random_proposal, gradient_rms) + / baseline_second_moment + ) + + observed_rho = float(score_gradient.statistic) + observed_ratio = information_second_moment / baseline_second_moment + return { + f"spearman_{score_label}_vs_gradient_rms": observed_rho, + f"spearman_{score_label}_vs_gradient_rms_pvalue_asymptotic": float( + score_gradient.pvalue + ), + f"kendall_{score_label}_vs_gradient_rms": float( + score_gradient_kendall.statistic + ), + f"kendall_{score_label}_vs_gradient_rms_pvalue_asymptotic": float( + score_gradient_kendall.pvalue + ), + f"spearman_{score_label}_vs_mean_loss": float(score_loss.statistic), + "top25_count": top_count, + "top25_overlap": len(score_top & gradient_top) / top_count, + f"top25_{score_label}_positions": [ + int(position) for position in sorted(score_top) + ], + "top25_gradient_positions": [ + int(position) for position in sorted(gradient_top) + ], + "proposal": { + "alpha": alpha, + "percentile_epsilon": percentile_epsilon, + "baseline_second_moment": baseline_second_moment, + "information_second_moment": information_second_moment, + "information_to_baseline_ratio": observed_ratio, + "optimal_second_moment": optimal_second_moment, + "optimal_to_baseline_ratio": ( + optimal_second_moment / baseline_second_moment + ), + "importance_weight_min": float(np.min(baseline / proposal)), + "importance_weight_max": float(np.max(baseline / proposal)), + }, + "random_permutation_control": { + "permutation_count": permutation_count, + "seed": permutation_seed, + "one_sided_pvalue_rho_at_least_observed": float( + (1 + np.count_nonzero(random_rhos >= observed_rho)) + / (permutation_count + 1) + ), + "one_sided_pvalue_second_moment_at_most_observed": float( + ( + 1 + + np.count_nonzero( + random_second_moment_ratios <= observed_ratio + ) + ) + / (permutation_count + 1) + ), + "second_moment_ratio_median": float( + np.median(random_second_moment_ratios) + ), + "second_moment_ratio_p05": float( + np.quantile(random_second_moment_ratios, 0.05) + ), + "second_moment_ratio_p95": float( + np.quantile(random_second_moment_ratios, 0.95) + ), + }, + } + + +def leave_one_out_robustness( + scores: np.ndarray, + gradient_rms: np.ndarray, + *, + alpha: float, + percentile_epsilon: float, +) -> dict: + """Check that alignment and variance gains survive every single omission.""" + scores = np.asarray(scores, dtype=np.float64) + gradient_rms = np.asarray(gradient_rms, dtype=np.float64) + if ( + scores.ndim != 1 + or gradient_rms.ndim != 1 + or scores.shape != gradient_rms.shape + or len(scores) < 4 + or not np.isfinite(scores).all() + or not np.isfinite(gradient_rms).all() + or np.any(gradient_rms <= 0.0) + ): + raise ValueError( + "scores and gradient_rms must be finite aligned arrays of length >= 4" + ) + + spearman_values = [] + second_moment_ratios = [] + for omitted in range(len(scores)): + keep = np.arange(len(scores)) != omitted + subset_scores = scores[keep] + subset_gradient_rms = gradient_rms[keep] + correlation = spearmanr(subset_scores, subset_gradient_rms).statistic + proposal, _ = proposal_from_scores( + subset_scores, + alpha=alpha, + percentile_epsilon=percentile_epsilon, + ) + baseline_second_moment = float(np.mean(np.square(subset_gradient_rms))) + ratio = ( + gradient_second_moment(proposal, subset_gradient_rms) + / baseline_second_moment + ) + spearman_values.append(float(correlation)) + second_moment_ratios.append(float(ratio)) + return { + "minimum_spearman": min(spearman_values), + "median_spearman": float(np.median(spearman_values)), + "maximum_second_moment_ratio": max(second_moment_ratios), + "median_second_moment_ratio": float(np.median(second_moment_ratios)), + "all_spearman_positive": all(value > 0.0 for value in spearman_values), + "all_second_moment_ratios_below_one": all( + value < 1.0 for value in second_moment_ratios + ), + } + + +def hierarchical_alignment_audit( + scores: np.ndarray, + gradient_rms: np.ndarray, + group_ids: np.ndarray, + *, + alpha: float, + percentile_epsilon: float, + permutation_count: int, + permutation_seed: int, +) -> dict: + """Audit proposals that only redistribute probability within each group. + + Groups represent episodes. Their marginal probability remains fixed; only + the conditional temporal proposal inside a group is changed. + """ + scores = np.asarray(scores, dtype=np.float64) + gradient_rms = np.asarray(gradient_rms, dtype=np.float64) + group_ids = np.asarray(group_ids) + if ( + scores.ndim != 1 + or gradient_rms.ndim != 1 + or group_ids.ndim != 1 + or not (scores.shape == gradient_rms.shape == group_ids.shape) + or len(scores) < 4 + ): + raise ValueError("hierarchical inputs must be aligned 1-D arrays") + if ( + not np.isfinite(scores).all() + or not np.isfinite(gradient_rms).all() + or np.any(gradient_rms <= 0.0) + ): + raise ValueError("scores must be finite and gradient_rms positive") + if permutation_count <= 0: + raise ValueError("permutation_count must be positive") + + unique_groups = list(dict.fromkeys(group_ids.tolist())) + if len(unique_groups) < 2: + raise ValueError("hierarchical audit requires at least two groups") + group_indices = [] + for group_id in unique_groups: + indices = np.flatnonzero(group_ids == group_id) + if len(indices) < 4: + raise ValueError("every hierarchical group needs at least four items") + group_indices.append(indices) + + group_weight = 1.0 / len(group_indices) + within_gradient_percentiles = np.empty(len(scores), dtype=np.float64) + for indices in group_indices: + within_gradient_percentiles[indices] = ( + rankdata(gradient_rms[indices], method="average") - 1.0 + ) / (len(indices) - 1.0) + + def evaluate(candidate_scores: np.ndarray) -> tuple[float, float, np.ndarray]: + information_moment = 0.0 + optimal_moment = 0.0 + within_percentiles = np.empty(len(scores), dtype=np.float64) + for indices in group_indices: + proposal, percentiles = proposal_from_scores( + candidate_scores[indices], + alpha=alpha, + percentile_epsilon=percentile_epsilon, + ) + group_gradient = gradient_rms[indices] + information_moment += group_weight * gradient_second_moment( + proposal, + group_gradient, + ) + optimal = group_gradient / group_gradient.sum() + optimal_moment += group_weight * gradient_second_moment( + optimal, + group_gradient, + ) + within_percentiles[indices] = percentiles + return information_moment, optimal_moment, within_percentiles + + baseline_moment = float( + sum( + group_weight * np.mean(np.square(gradient_rms[indices])) + for indices in group_indices + ) + ) + information_moment, optimal_moment, within_percentiles = evaluate(scores) + observed_ratio = information_moment / baseline_moment + pooled_correlation = spearmanr( + within_percentiles, + within_gradient_percentiles, + ) + + per_group = [] + for group_id, indices in zip(unique_groups, group_indices): + proposal, percentiles = proposal_from_scores( + scores[indices], + alpha=alpha, + percentile_epsilon=percentile_epsilon, + ) + group_gradient = gradient_rms[indices] + baseline = float(np.mean(np.square(group_gradient))) + information = gradient_second_moment(proposal, group_gradient) + correlation = spearmanr(percentiles, group_gradient) + top_count = max(1, len(indices) // 4) + score_top = set(np.argsort(percentiles)[-top_count:]) + gradient_top = set(np.argsort(group_gradient)[-top_count:]) + per_group.append( + { + "group_id": str(group_id), + "sample_count": len(indices), + "spearman": float(correlation.statistic), + "top25_overlap": len(score_top & gradient_top) / top_count, + "baseline_second_moment": baseline, + "information_second_moment": information, + "information_to_baseline_ratio": information / baseline, + } + ) + + rng = np.random.default_rng(permutation_seed) + random_rhos = np.empty(permutation_count, dtype=np.float64) + random_ratios = np.empty(permutation_count, dtype=np.float64) + for permutation_index in range(permutation_count): + shuffled_scores = scores.copy() + for indices in group_indices: + shuffled_scores[indices] = rng.permutation( + shuffled_scores[indices] + ) + random_moment, _, random_percentiles = evaluate(shuffled_scores) + random_ratios[permutation_index] = random_moment / baseline_moment + random_rhos[permutation_index] = spearmanr( + random_percentiles, + within_gradient_percentiles, + ).statistic + + group_ratios = [ + group["information_to_baseline_ratio"] for group in per_group + ] + leave_one_group_out_ratios = [] + for omitted in range(len(per_group)): + kept = [ + group for index, group in enumerate(per_group) + if index != omitted + ] + kept_baseline = np.mean( + [group["baseline_second_moment"] for group in kept] + ) + kept_information = np.mean( + [group["information_second_moment"] for group in kept] + ) + leave_one_group_out_ratios.append( + float(kept_information / kept_baseline) + ) + + return { + "group_count": len(group_indices), + "group_marginal": "fixed_equal_empirical_weight", + "pooled_within_group_spearman": float( + pooled_correlation.statistic + ), + "pooled_within_group_spearman_pvalue_asymptotic": float( + pooled_correlation.pvalue + ), + "median_group_spearman": float( + np.median([group["spearman"] for group in per_group]) + ), + "mean_group_top25_overlap": float( + np.mean([group["top25_overlap"] for group in per_group]) + ), + "groups_with_positive_spearman": sum( + group["spearman"] > 0.0 for group in per_group + ), + "groups_with_second_moment_improvement": sum( + ratio < 1.0 for ratio in group_ratios + ), + "proposal": { + "alpha": alpha, + "percentile_epsilon": percentile_epsilon, + "baseline_second_moment": baseline_moment, + "information_second_moment": information_moment, + "information_to_baseline_ratio": observed_ratio, + "optimal_second_moment": optimal_moment, + "optimal_to_baseline_ratio": optimal_moment / baseline_moment, + }, + "random_permutation_control": { + "permutation_count": permutation_count, + "seed": permutation_seed, + "one_sided_pvalue_rho_at_least_observed": float( + (1 + np.count_nonzero( + random_rhos >= pooled_correlation.statistic + )) + / (permutation_count + 1) + ), + "one_sided_pvalue_second_moment_at_most_observed": float( + (1 + np.count_nonzero(random_ratios <= observed_ratio)) + / (permutation_count + 1) + ), + "second_moment_ratio_median": float(np.median(random_ratios)), + "second_moment_ratio_p05": float( + np.quantile(random_ratios, 0.05) + ), + "second_moment_ratio_p95": float( + np.quantile(random_ratios, 0.95) + ), + }, + "robustness": { + "maximum_leave_one_group_out_ratio": max( + leave_one_group_out_ratios + ), + "all_leave_one_group_out_ratios_below_one": all( + ratio < 1.0 for ratio in leave_one_group_out_ratios + ), + }, + "per_group": per_group, + } diff --git a/real_robot/source/scripts/build_residual_temporal_proposal.py b/real_robot/source/scripts/build_residual_temporal_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..0792a90b5f8ddf0f4f55687fe6f0fe3d2696ff18 --- /dev/null +++ b/real_robot/source/scripts/build_residual_temporal_proposal.py @@ -0,0 +1,10 @@ +"""Compatibility entrypoint for the ReTime proposal builder.""" +from retime.commands import proposal as _implementation + + +def __getattr__(name): + return getattr(_implementation, name) + + +if __name__ == "__main__": + _implementation.main() diff --git a/real_robot/source/scripts/build_shuffle_temporal_proposal.py b/real_robot/source/scripts/build_shuffle_temporal_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..2c1a5adce91c391904419f9b2450d8a2f0d02cac --- /dev/null +++ b/real_robot/source/scripts/build_shuffle_temporal_proposal.py @@ -0,0 +1,10 @@ +"""Compatibility entrypoint; implementation lives in the ReTime plugin.""" +from retime.commands import shuffle as _implementation + + +def __getattr__(name): + return getattr(_implementation, name) + + +if __name__ == "__main__": + _implementation.main() \ No newline at end of file diff --git a/real_robot/source/scripts/convert_xtrainer2_hdf5_to_lerobot.py b/real_robot/source/scripts/convert_xtrainer2_hdf5_to_lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..8fcbbeb2b5520aac13c41d784fedb4231edd317b --- /dev/null +++ b/real_robot/source/scripts/convert_xtrainer2_hdf5_to_lerobot.py @@ -0,0 +1,722 @@ +#!/usr/bin/env python3 +"""Convert xtrainer2 optimized HDF5 episodes to StarVLA-compatible LeRobot v2. + +The source streams run at different rates. Every modality therefore uses its +own ``aligned_index10`` array; indexing all streams with the control-frame +number directly silently corrupts alignment. + +Rotations are stored as ``x, y, z, w`` in the source but as ``w, x, y, z`` in +the output, because StarVLA's rotation transforms follow the PyTorch3D +convention. Anything that replays these actions on the robot must undo the +reordering. + +Training language is canonical English. The original Chinese instruction is +kept as ``source_instruction`` in episode metadata. +""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import hashlib +import io +import json +import os +import shutil +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable + +import h5py +import numpy as np +import pandas as pd +from PIL import Image + + +FPS = 10 +CHUNK_SIZE = 1000 +CAMERAS = ("cam_high", "cam_left_wrist", "cam_right_wrist") +STATE_KEY = "observation.state" +ACTION_KEY = "action" +QUATERNION_STARTS = (3, 11) +SOURCE_QUATERNION_FORMAT = "xyzw" +OUTPUT_QUATERNION_FORMAT = "wxyz" +DEFAULT_INPUT = Path("/apdcephfs_nj7/share_305204761/arianliu/data/benchmark") +DEFAULT_OUTPUT = Path("/apdcephfs_nj7/share_305204761/arianliu/data/lerobot") +TASK_SOURCES = { + "xtrainer2_pick_bottle": ("pick_bottle_dataset_*",), + "xtrainer2_pack_phone": ("pack_phone_dataset_*",), +} +CANONICAL_ENGLISH = { + "夹取粉色妆前液瓶放入方碟中": "Pick up the pink primer bottle and place it in the square dish.", + "夹取粉色滴管放入方碟中": "Pick up the pink dropper and place it in the square dish.", + "夹取紫色眉胶管放入方碟中": "Pick up the purple eyebrow gel tube and place it in the square dish.", + "夹取银色眉胶管放入方碟中": "Pick up the silver eyebrow gel tube and place it in the square dish.", + "打包手机": "Pick up the phone, place it in the box, and close the lid.", +} + + +@dataclass(frozen=True) +class SourceEpisode: + path: Path + episode_index: int + length: int + instruction: str + source_instruction: str + task_index: int + source_id: str + global_start: int + + +def _json_dump(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + with tmp.open("w", encoding="utf-8") as handle: + json.dump(value, handle, ensure_ascii=False, indent=2) + handle.write("\n") + os.replace(tmp, path) + + +def _jsonl_dump(path: Path, rows: Iterable[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(path.suffix + ".tmp") + with tmp.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row, ensure_ascii=False) + "\n") + os.replace(tmp, path) + + +def _decode_text(value: Any) -> str: + if isinstance(value, bytes): + return value.decode("utf-8") + return str(value) + + +def canonicalize_instruction(source_instruction: str) -> str: + text = source_instruction.strip() + if not text: + raise ValueError("Empty instruction") + if text in CANONICAL_ENGLISH: + return CANONICAL_ENGLISH[text] + if text in CANONICAL_ENGLISH.values(): + return text + raise ValueError(f"Unknown instruction {text!r}") + + +def discover_sources(input_root: Path, dataset_name: str) -> list[SourceEpisode]: + patterns = TASK_SOURCES[dataset_name] + paths: list[Path] = [] + for pattern in patterns: + for directory in sorted(input_root.glob(pattern)): + paths.extend(sorted(directory.glob("*_optimized.hdf5"))) + if not paths: + raise FileNotFoundError(f"No source HDF5 files found for {dataset_name} below {input_root}") + + source_rows: list[tuple[Path, int, str, str]] = [] + global_start = 0 + seen_ids: set[str] = set() + for episode_index, path in enumerate(paths): + with h5py.File(path, "r") as source: + length = int(source.attrs["n_frames_10"]) + instruction = _decode_text(source.attrs["instruction"]).strip() + source_id = _decode_text(source.attrs.get("id", path.stem)) + if source_id in seen_ids: + raise ValueError(f"Duplicate source id {source_id!r}: {path}") + if length <= 0 or not instruction: + raise ValueError(f"Invalid episode metadata in {path}: length={length}, instruction={instruction!r}") + canonicalize_instruction(instruction) + seen_ids.add(source_id) + source_rows.append((path, length, instruction, source_id)) + task_indices = { + source_instruction: index + for index, source_instruction in enumerate(sorted({row[2] for row in source_rows})) + } + episodes: list[SourceEpisode] = [] + for episode_index, (path, length, source_instruction, source_id) in enumerate(source_rows): + episodes.append( + SourceEpisode( + path, + episode_index, + length, + canonicalize_instruction(source_instruction), + source_instruction, + task_indices[source_instruction], + source_id, + global_start, + ) + ) + global_start += length + return episodes + + +def _aligned(source: h5py.File, group: str, expected_length: int) -> np.ndarray: + indices = np.asarray(source[f"{group}/aligned_index10"], dtype=np.int64) + values = source[f"{group}/value"] + if indices.shape != (expected_length,): + raise ValueError(f"{group}: aligned_index10 has shape {indices.shape}, expected {(expected_length,)}") + if indices.min(initial=0) < 0 or indices.max(initial=-1) >= len(values): + raise ValueError(f"{group}: aligned indices outside [0, {len(values)})") + if np.any(np.diff(indices) < 0): + raise ValueError(f"{group}: aligned indices are not monotonic") + return indices + + +def xyzw_to_wxyz(values: np.ndarray) -> np.ndarray: + """Rewrite both quaternion blocks of a 16-D vector series into wxyz order.""" + converted = np.array(values, dtype=np.float32, copy=True) + for start in QUATERNION_STARTS: + block = values[:, start : start + 4] + converted[:, start : start + 4] = block[:, [3, 0, 1, 2]] + return converted + + +def load_low_dim(source: h5py.File, length: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + action_idx = _aligned(source, "action", length) + state_idx = _aligned(source, "observations/qpos", length) + actions = np.asarray(source["action/value"][action_idx], dtype=np.float32) + states = np.asarray(source["observations/qpos/value"][state_idx], dtype=np.float32) + timestamps = np.asarray(source["aligned_timestamp10"], dtype=np.float64) + if actions.shape != (length, 16) or states.shape != (length, 16): + raise ValueError(f"Expected 16-D action/state, got {actions.shape} and {states.shape}") + if timestamps.shape != (length,) or np.any(np.diff(timestamps) <= 0): + raise ValueError("aligned_timestamp10 must be strictly increasing") + for name, values in (("action", actions), ("state", states)): + if not np.isfinite(values).all(): + raise ValueError(f"{name} contains non-finite values") + for start in QUATERNION_STARTS: + norms = np.linalg.norm(values[:, start : start + 4], axis=-1) + if np.any(np.abs(norms - 1.0) > 5e-2): + raise ValueError(f"{name} contains non-unit quaternions (max error {np.max(np.abs(norms - 1)):.4g})") + return xyzw_to_wxyz(states), xyzw_to_wxyz(actions), timestamps + + +def _probe_jpeg(encoded: np.ndarray, camera: str) -> tuple[int, int]: + payload = np.asarray(encoded, dtype=np.uint8).tobytes() + if not payload.startswith(b"\xff\xd8"): + raise ValueError(f"{camera}: frame is not JPEG") + with Image.open(io.BytesIO(payload)) as image: + image.verify() + return image.size + + +def _ffmpeg_command(ffmpeg: str, output: Path) -> list[str]: + return [ + ffmpeg, + "-hide_banner", + "-loglevel", + "error", + "-y", + "-f", + "image2pipe", + "-vcodec", + "mjpeg", + "-framerate", + str(FPS), + "-i", + "pipe:0", + "-an", + "-c:v", + "libopenh264", + "-b:v", + "4M", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + "-f", + "mp4", + str(output), + ] + + +def write_video( + source: h5py.File, + camera: str, + length: int, + output: Path, + ffmpeg: str, +) -> tuple[int, int]: + group = f"observations/images/{camera}" + indices = _aligned(source, group, length) + values = source[f"{group}/value"] + first_size = _probe_jpeg(values[int(indices[0])], camera) + output.parent.mkdir(parents=True, exist_ok=True) + partial = output.with_suffix(".part.mp4") + partial.unlink(missing_ok=True) + process = subprocess.Popen( + _ffmpeg_command(ffmpeg, partial), + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + ) + assert process.stdin is not None + try: + for source_index in indices: + process.stdin.write(np.asarray(values[int(source_index)], dtype=np.uint8).tobytes()) + process.stdin.close() + stderr = process.stderr.read().decode("utf-8", errors="replace") if process.stderr else "" + return_code = process.wait() + except Exception: + process.kill() + process.wait() + partial.unlink(missing_ok=True) + raise + if return_code: + partial.unlink(missing_ok=True) + raise RuntimeError(f"ffmpeg failed for {camera}: {stderr.strip()}") + if not partial.exists() or partial.stat().st_size == 0: + raise RuntimeError(f"ffmpeg produced an empty video for {camera}") + os.replace(partial, output) + return first_size + + +def _vector_series(values: np.ndarray) -> list[list[float]]: + return [row.tolist() for row in np.asarray(values, dtype=np.float32)] + + +def write_parquet( + output: Path, + episode: SourceEpisode, + states: np.ndarray, + actions: np.ndarray, +) -> None: + length = episode.length + frame_index = np.arange(length, dtype=np.int64) + table = pd.DataFrame( + { + STATE_KEY: _vector_series(states), + ACTION_KEY: _vector_series(actions), + "timestamp": frame_index.astype(np.float64) / FPS, + "frame_index": frame_index, + "episode_index": np.full(length, episode.episode_index, dtype=np.int64), + "index": np.arange(episode.global_start, episode.global_start + length, dtype=np.int64), + "task_index": np.full(length, episode.task_index, dtype=np.int64), + } + ) + output.parent.mkdir(parents=True, exist_ok=True) + partial = output.with_suffix(".part.parquet") + partial.unlink(missing_ok=True) + table.to_parquet(partial, engine="pyarrow", index=False) + os.replace(partial, output) + + +def chunk_dir(episode_index: int) -> str: + return f"chunk-{episode_index // CHUNK_SIZE:03d}" + + +def parquet_path(root: Path, episode_index: int) -> Path: + return root / "data" / chunk_dir(episode_index) / f"episode_{episode_index:06d}.parquet" + + +def video_path(root: Path, episode_index: int, camera: str) -> Path: + return ( + root + / "videos" + / chunk_dir(episode_index) + / f"observation.images.{camera}" + / f"episode_{episode_index:06d}.mp4" + ) + + +def episode_complete(root: Path, episode: SourceEpisode) -> bool: + parquet = parquet_path(root, episode.episode_index) + videos = [video_path(root, episode.episode_index, camera) for camera in CAMERAS] + if not parquet.exists() or any(not path.exists() for path in videos): + return False + try: + return len(pd.read_parquet(parquet, columns=["frame_index"])) == episode.length + except Exception: + return False + + +def convert_episode( + root: Path, + episode: SourceEpisode, + ffmpeg: str, + overwrite: bool, +) -> dict[str, Any]: + if episode_complete(root, episode) and not overwrite: + return {"status": "skipped", "episode_index": episode.episode_index} + targets = [parquet_path(root, episode.episode_index)] + targets.extend(video_path(root, episode.episode_index, camera) for camera in CAMERAS) + if overwrite: + for target in targets: + target.unlink(missing_ok=True) + + with h5py.File(episode.path, "r") as source: + states, actions, _ = load_low_dim(source, episode.length) + sizes = {} + for camera in CAMERAS: + sizes[camera] = write_video( + source, + camera, + episode.length, + video_path(root, episode.episode_index, camera), + ffmpeg, + ) + write_parquet(parquet_path(root, episode.episode_index), episode, states, actions) + return { + "status": "converted", + "episode_index": episode.episode_index, + "length": episode.length, + "source": str(episode.path), + "image_sizes": {key: list(value) for key, value in sizes.items()}, + } + + +def modality_metadata() -> dict[str, Any]: + arm_fields = { + "left_position": {"start": 0, "end": 3, "dtype": "float32", "absolute": True}, + "left_rotation": { + "start": 3, + "end": 7, + "dtype": "float32", + "absolute": True, + "rotation_type": "quaternion", + }, + "left_gripper": { + "start": 7, + "end": 8, + "dtype": "float32", + "absolute": True, + }, + "right_position": {"start": 8, "end": 11, "dtype": "float32", "absolute": True}, + "right_rotation": { + "start": 11, + "end": 15, + "dtype": "float32", + "absolute": True, + "rotation_type": "quaternion", + }, + "right_gripper": { + "start": 15, + "end": 16, + "dtype": "float32", + "absolute": True, + }, + } + state = {key: {**value, "original_key": STATE_KEY} for key, value in arm_fields.items()} + action = {key: {**value, "original_key": ACTION_KEY} for key, value in arm_fields.items()} + return { + "state": state, + "action": action, + "video": { + camera: {"original_key": f"observation.images.{camera}"} for camera in CAMERAS + }, + "annotation": { + "human.action.task_description": {"original_key": "task_index"}, + }, + } + + +def _statistics(values: np.ndarray) -> dict[str, list[float]]: + values = np.asarray(values, dtype=np.float64) + return { + "max": np.max(values, axis=0).tolist(), + "min": np.min(values, axis=0).tolist(), + "mean": np.mean(values, axis=0).tolist(), + "std": np.std(values, axis=0).tolist(), + "q01": np.quantile(values, 0.01, axis=0).tolist(), + "q99": np.quantile(values, 0.99, axis=0).tolist(), + } + + +def build_stats(episodes: list[SourceEpisode]) -> dict[str, Any]: + all_states, all_actions = [], [] + for episode in episodes: + with h5py.File(episode.path, "r") as source: + states, actions, _ = load_low_dim(source, episode.length) + all_states.append(states) + all_actions.append(actions) + statistics = { + STATE_KEY: _statistics(np.concatenate(all_states, axis=0)), + ACTION_KEY: _statistics(np.concatenate(all_actions, axis=0)), + } + return { + "__format_version": 2, + "__cache_config": {"mode": "abs"}, + "statistics": statistics, + } + + +def build_info( + dataset_name: str, + episodes: list[SourceEpisode], + image_size: tuple[int, int], +) -> dict[str, Any]: + width, height = image_size + total_frames = sum(episode.length for episode in episodes) + features: dict[str, Any] = { + STATE_KEY: {"dtype": "float32", "shape": [16], "names": None}, + ACTION_KEY: {"dtype": "float32", "shape": [16], "names": None}, + } + for camera in CAMERAS: + features[f"observation.images.{camera}"] = { + "dtype": "video", + "shape": [height, width, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": FPS, + "video.codec": "h264", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + } + return { + "codebase_version": "v2.0", + "robot_type": "xtrainer2_bimanual_eef", + "fps": FPS, + "total_episodes": len(episodes), + "total_frames": total_frames, + "total_tasks": len({episode.instruction for episode in episodes}), + "total_videos": len(episodes) * len(CAMERAS), + "total_chunks": (len(episodes) + CHUNK_SIZE - 1) // CHUNK_SIZE, + "chunks_size": CHUNK_SIZE, + "splits": {"train": f"0:{len(episodes)}"}, + "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4", + "features": features, + "dataset_name": dataset_name, + "quaternion_format": OUTPUT_QUATERNION_FORMAT, + "source_quaternion_format": SOURCE_QUATERNION_FORMAT, + } + + +def source_fingerprint(episodes: list[SourceEpisode]) -> str: + digest = hashlib.sha256() + for episode in episodes: + stat = episode.path.stat() + digest.update( + f"{episode.path}:{stat.st_size}:{stat.st_mtime_ns}:{episode.length}:{episode.source_id}\n".encode() + ) + return digest.hexdigest() + + +def write_metadata( + root: Path, + dataset_name: str, + episodes: list[SourceEpisode], + image_size: tuple[int, int], +) -> None: + tasks = sorted( + {(episode.task_index, episode.instruction) for episode in episodes} + ) + _json_dump(root / "meta" / "modality.json", modality_metadata()) + _json_dump(root / "meta" / "info.json", build_info(dataset_name, episodes, image_size)) + _json_dump(root / "meta" / "stats_gr00t.json", build_stats(episodes)) + _jsonl_dump( + root / "meta" / "episodes.jsonl", + ( + { + "episode_index": episode.episode_index, + "tasks": [episode.instruction], + "source_instruction": episode.source_instruction, + "length": episode.length, + "source_file": str(episode.path), + "source_id": episode.source_id, + } + for episode in episodes + ), + ) + _jsonl_dump( + root / "meta" / "tasks.jsonl", + ({"task_index": task_index, "task": instruction} for task_index, instruction in tasks), + ) + + +def _jsonl_load(path: Path) -> list[dict[str, Any]]: + with path.open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def relabel_dataset_language(root: Path) -> dict[str, Any]: + """Rewrite training language to canonical English without reconverting videos.""" + tasks_path = root / "meta" / "tasks.jsonl" + episodes_path = root / "meta" / "episodes.jsonl" + if not tasks_path.exists() or not episodes_path.exists(): + raise FileNotFoundError(f"Language metadata missing under {root}") + + english_to_source = {english: chinese for chinese, english in CANONICAL_ENGLISH.items()} + + def split_language(raw: str) -> tuple[str, str]: + text = raw.strip() + if text in CANONICAL_ENGLISH: + return CANONICAL_ENGLISH[text], text + if text in english_to_source: + return text, english_to_source[text] + raise ValueError(f"Unknown instruction {text!r} in {root}") + + rewritten_tasks = [] + for row in _jsonl_load(tasks_path): + english, source = split_language(row.get("source_instruction") or row["task"]) + rewritten_tasks.append({"task_index": row["task_index"], "task": english}) + + rewritten_episodes = [] + for row in _jsonl_load(episodes_path): + english, source = split_language(row.get("source_instruction") or row["tasks"][0]) + rewritten_episodes.append({**row, "tasks": [english], "source_instruction": source}) + + _jsonl_dump(tasks_path, rewritten_tasks) + _jsonl_dump(episodes_path, rewritten_episodes) + return { + "dataset": root.name, + "tasks": rewritten_tasks, + "episodes": len(rewritten_episodes), + } + + +def find_ffmpeg(requested: str | None) -> str: + candidate = requested or shutil.which("ffmpeg") + if not candidate: + conda_ffmpeg = Path("/opt/conda/envs/torch-base/bin/ffmpeg") + candidate = str(conda_ffmpeg) if conda_ffmpeg.exists() else None + if not candidate: + raise FileNotFoundError("ffmpeg was not found; pass --ffmpeg /path/to/ffmpeg") + return candidate + + +def convert_dataset(args: argparse.Namespace, dataset_name: str) -> dict[str, Any]: + episodes = discover_sources(args.input_root, dataset_name) + if args.max_episodes is not None: + episodes = episodes[: args.max_episodes] + root = args.output_root / dataset_name + root.mkdir(parents=True, exist_ok=True) + fingerprint = source_fingerprint(episodes) + state_path = root / "meta" / "conversion_state.json" + if state_path.exists() and not args.overwrite: + with state_path.open(encoding="utf-8") as handle: + old_state = json.load(handle) + if old_state.get("source_fingerprint") != fingerprint: + raise RuntimeError( + f"Source set changed for {dataset_name}; use --overwrite or a different output directory" + ) + + ffmpeg = find_ffmpeg(args.ffmpeg) + started = time.time() + results, errors = [], [] + image_size: tuple[int, int] | None = None + + def run_one(episode: SourceEpisode): + try: + return episode, convert_episode(root, episode, ffmpeg, args.overwrite), None + except Exception as exc: + return episode, None, exc + + with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor: + converted = executor.map(run_one, episodes) + for position, (episode, result, exc) in enumerate(converted, start=1): + if exc is None: + assert result is not None + results.append(result) + if result.get("image_sizes"): + sizes = {tuple(value) for value in result["image_sizes"].values()} + if len(sizes) != 1: + exc = ValueError(f"Camera resolutions differ: {result['image_sizes']}") + else: + current_size = next(iter(sizes)) + if image_size is not None and current_size != image_size: + exc = ValueError( + f"Image resolution changed from {image_size} to {current_size}" + ) + image_size = current_size + if exc is None: + assert result is not None + print( + f"[{dataset_name}] {position}/{len(episodes)} episode {episode.episode_index}: " + f"{result['status']}", + flush=True, + ) + else: + error = { + "episode_index": episode.episode_index, + "source": str(episode.path), + "error": f"{type(exc).__name__}: {exc}", + } + errors.append(error) + print( + f"[{dataset_name}] ERROR {json.dumps(error, ensure_ascii=False)}", + file=sys.stderr, + flush=True, + ) + _json_dump( + state_path, + { + "dataset_name": dataset_name, + "source_fingerprint": fingerprint, + "expected_episodes": len(episodes), + "results": results, + "errors": errors, + "updated_at": time.time(), + }, + ) + + if image_size is None: + with h5py.File(episodes[0].path, "r") as source: + image_size = _probe_jpeg(source[f"observations/images/{CAMERAS[0]}/value"][0], CAMERAS[0]) + complete = not errors and all(episode_complete(root, episode) for episode in episodes) + if complete: + write_metadata(root, dataset_name, episodes, image_size) + report = { + "dataset_name": dataset_name, + "complete": complete, + "output_root": str(root), + "source_fingerprint": fingerprint, + "episodes_expected": len(episodes), + "episodes_complete": sum(episode_complete(root, episode) for episode in episodes), + "total_frames": sum(episode.length for episode in episodes), + "image_size": list(image_size), + "elapsed_seconds": time.time() - started, + "errors": errors, + } + _json_dump(root / "conversion_report.json", report) + return report + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input-root", type=Path, default=DEFAULT_INPUT) + parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument( + "--datasets", + nargs="+", + choices=sorted(TASK_SOURCES), + default=sorted(TASK_SOURCES), + ) + parser.add_argument("--max-episodes", type=int, default=None) + parser.add_argument("--ffmpeg", type=str, default=None) + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--keep-going", action="store_true") + parser.add_argument( + "--relabel-language", + action="store_true", + help="Rewrite tasks.jsonl and episodes.jsonl to canonical English without reconverting videos", + ) + args = parser.parse_args() + if args.max_episodes is not None and args.max_episodes <= 0: + parser.error("--max-episodes must be positive") + if args.workers <= 0: + parser.error("--workers must be positive") + return args + + +def main() -> int: + args = parse_args() + if args.relabel_language: + reports = [ + relabel_dataset_language(args.output_root / dataset_name) + for dataset_name in args.datasets + ] + print(json.dumps(reports, ensure_ascii=False, indent=2)) + return 0 + reports = [convert_dataset(args, dataset_name) for dataset_name in args.datasets] + print(json.dumps(reports, ensure_ascii=False, indent=2)) + return 0 if all(report["complete"] for report in reports) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/real_robot/source/scripts/launch_pack_phone_optr_proposal.sh b/real_robot/source/scripts/launch_pack_phone_optr_proposal.sh new file mode 100644 index 0000000000000000000000000000000000000000..68ff2a1c57eb438bfd5179fedde0a7a080c49516 --- /dev/null +++ b/real_robot/source/scripts/launch_pack_phone_optr_proposal.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +# Forward arguments to the ReTime command-line interface. +set -euo pipefail +project="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +export PYTHONPATH="${project}:${PYTHONPATH:-}" +exec "${OPTR_PYTHON:-python}" -m retime "$@" diff --git a/real_robot/source/scripts/launch_pack_phone_residual_20ep_to_60ep.sh b/real_robot/source/scripts/launch_pack_phone_residual_20ep_to_60ep.sh new file mode 100644 index 0000000000000000000000000000000000000000..a676cbe154bb179b7c8370f61f7b56e392b5fdc2 --- /dev/null +++ b/real_robot/source/scripts/launch_pack_phone_residual_20ep_to_60ep.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +export OPTR_PROJECT_ROOT="$project" +export OPTR_PYTHON="${OPTR_PYTHON:-/opt/conda/envs/torch-base/bin/python}" +export OPTR_RUN_ROOT="${OPTR_RUN_ROOT:-/path/to/runs}" +export FLORENCE_MODEL_PATH="${FLORENCE_MODEL_PATH:-/path/to/models/Florence-2-large}" +export XTRAINER2_DATA_ROOT="${XTRAINER2_DATA_ROOT:-/path/to/data/lerobot}" +export XTRAINER2_PACK_PHONE_PROVENANCE_PATH="${XTRAINER2_PACK_PHONE_PROVENANCE_PATH:-${XTRAINER2_DATA_ROOT}/xtrainer2_optr_support/pack_phone/data_provenance.json}" +export XTRAINER2_PACK_PHONE_FINGERPRINT_PATH="${XTRAINER2_PACK_PHONE_FINGERPRINT_PATH:-${XTRAINER2_DATA_ROOT}/xtrainer2_optr_support/pack_phone/content_fingerprints.json}" + +parent_run="${OPTR_PACK_PHONE_PARENT_RUN:-${OPTR_RUN_ROOT}/florence-xtrainer2-pack-phone-uniform}" +export XTRAINER2_PACK_PHONE_RESUME_DIR="${XTRAINER2_PACK_PHONE_RESUME_DIR:-${parent_run}/checkpoints/steps_17460_accelerate}" +export XTRAINER2_PACK_PHONE_MANIFEST_PATH="${XTRAINER2_PACK_PHONE_MANIFEST_PATH:-${parent_run}/valid_anchor_manifest.json}" +export OPTR_ACCELERATE_CONFIG="${OPTR_ACCELERATE_CONFIG:-${project}/configs/deepspeed_florence_8gpu.yaml}" +export OPTR_NUM_PROCESSES=8 +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_RUN_ID="${WANDB_RUN_ID:-pack-phone-residual-optr-20ep-to-60ep}" +export WANDB_RESUME=allow + +: "${XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH:?set the locked common-start integrity manifest}" +: "${OPTR_TEMPORAL_PROPOSAL_PATH:?set the residual proposal JSON}" + +run_id="${OPTR_RUN_ID:-florence-xtrainer2-pack-phone-residual-20ep-to-60ep}" +log="${OPTR_RUN_ROOT}/${run_id}.log" +mkdir -p "$OPTR_RUN_ROOT" + +echo "launching exact residual continuation: $run_id" +echo " common-start: $XTRAINER2_PACK_PHONE_RESUME_DIR" +echo " proposal: $OPTR_TEMPORAL_PROPOSAL_PATH" +echo " log: $log" +exec "${project}/scripts/run_xtrainer2_arm.sh" \ + residual "$run_id" "$OPTR_TEMPORAL_PROPOSAL_PATH" \ + "${project}/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml" \ + >>"$log" 2>&1 diff --git a/real_robot/source/scripts/launch_pack_phone_uniform.sh b/real_robot/source/scripts/launch_pack_phone_uniform.sh new file mode 100644 index 0000000000000000000000000000000000000000..96005708f41ee1b69402452e2f5d5c5239198560 --- /dev/null +++ b/real_robot/source/scripts/launch_pack_phone_uniform.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# Launch the 60-epoch pack_phone uniform run. +# +# Global batch 128 (per-device 4 x 8 GPUs x accumulation 4), 874 optimizer +# steps per data epoch, 52380 steps total, checkpoints every 10 epochs at +# [8730, 17460, 26190, 34920, 43650, 52380]. Roughly 80 hours at the measured +# ~5.4 s per optimizer step. +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +export OPTR_PROJECT_ROOT=$PWD +export OPTR_PYTHON="${OPTR_PYTHON:-python}" +export OPTR_RUN_ROOT=/path/to/runs +export FLORENCE_MODEL_PATH=/path/to/models/Florence-2-large +export XTRAINER2_DATA_ROOT=/path/to/data/lerobot +export XTRAINER2_PACK_PHONE_PROVENANCE_PATH=$XTRAINER2_DATA_ROOT/xtrainer2_optr_support/pack_phone/data_provenance.json +export XTRAINER2_PACK_PHONE_FINGERPRINT_PATH=$XTRAINER2_DATA_ROOT/xtrainer2_optr_support/pack_phone/content_fingerprints.json +export OPTR_ACCELERATE_CONFIG=$OPTR_PROJECT_ROOT/configs/deepspeed_florence_8gpu.yaml +export OPTR_NUM_PROCESSES=8 + +# Stream metrics to wandb.ai. Requires the proxy for the whole run, not just +# the initial model read, and a key in ~/.netrc (wandb >= 0.30 -- 0.21.x +# rejects the newer wandb_v1_ key format as "must be 40 characters"). +# Set WANDB_MODE=offline to fall back to local files plus a later `wandb sync`. +export WANDB_MODE="${WANDB_MODE:-offline}" + +# Resume the same wandb run if the job is relaunched, so a restart continues +# one set of curves instead of starting a second run. +export WANDB_RUN_ID="${WANDB_RUN_ID:-pack-phone-uniform-60ep}" +export WANDB_RESUME=allow + +run_id=florence-xtrainer2-pack-phone-uniform +log=$OPTR_RUN_ROOT/${run_id}.log +mkdir -p "$OPTR_RUN_ROOT" + +echo "launching ${run_id}" +echo " log: ${log}" +exec bash examples/Xtrainer2/train_files/run_xtrainer2_train.sh \ + configs/florence_xtrainer2_pack_phone_uniform.yaml \ + "$run_id" >> "$log" 2>&1 diff --git a/real_robot/source/scripts/launch_pick_bottle_rel_aug_smoke.sh b/real_robot/source/scripts/launch_pick_bottle_rel_aug_smoke.sh new file mode 100644 index 0000000000000000000000000000000000000000..4537db4dcc2fca2e11eec611396632281880d575 --- /dev/null +++ b/real_robot/source/scripts/launch_pick_bottle_rel_aug_smoke.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Short ZeRO-3 smoke for pick_bottle rel-EE + image aug (6 steps, save at 4). +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +export OPTR_PROJECT_ROOT=$PWD +export OPTR_PYTHON="${OPTR_PYTHON:-python}" +export OPTR_RUN_ROOT=/path/to/runs +export FLORENCE_MODEL_PATH=/path/to/models/Florence-2-large +export XTRAINER2_DATA_ROOT=/path/to/data/lerobot +export XTRAINER2_PICK_BOTTLE_PROVENANCE_PATH=$XTRAINER2_DATA_ROOT/xtrainer2_optr_support/pick_bottle/data_provenance.json +export XTRAINER2_PICK_BOTTLE_FINGERPRINT_PATH=$XTRAINER2_DATA_ROOT/xtrainer2_optr_support/pick_bottle/content_fingerprints.json +export OPTR_ACCELERATE_CONFIG=$OPTR_PROJECT_ROOT/configs/deepspeed_florence_8gpu.yaml +export OPTR_NUM_PROCESSES=8 + +export WANDB_MODE="${WANDB_MODE:-offline}" + +run_id=florence-xtrainer2-pick-bottle-rel-aug-smoke +log=$OPTR_RUN_ROOT/${run_id}.log +mkdir -p "$OPTR_RUN_ROOT" + +echo "launching ${run_id}" +echo " log: ${log}" +exec bash examples/Xtrainer2/train_files/run_xtrainer2_train.sh \ + configs/florence_xtrainer2_pick_bottle_rel_aug_smoke.yaml \ + "$run_id" >> "$log" 2>&1 diff --git a/real_robot/source/scripts/launch_pick_bottle_rel_aug_uniform.sh b/real_robot/source/scripts/launch_pick_bottle_rel_aug_uniform.sh new file mode 100644 index 0000000000000000000000000000000000000000..f4353189935d48b471524a23a6b232b04ec65c88 --- /dev/null +++ b/real_robot/source/scripts/launch_pick_bottle_rel_aug_uniform.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Launch the 60-epoch pick_bottle uniform run (relative EE + image aug). +# +# Global batch 128 (per-device 4 x 8 GPUs x accumulation 4), 592 optimizer +# steps per data epoch, 35520 steps total, checkpoints every 10 epochs at +# [5920, 11840, 17760, 23680, 29600, 35520]. +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +export OPTR_PROJECT_ROOT=$PWD +export OPTR_PYTHON="${OPTR_PYTHON:-python}" +export OPTR_RUN_ROOT=/path/to/runs +export FLORENCE_MODEL_PATH=/path/to/models/Florence-2-large +export XTRAINER2_DATA_ROOT=/path/to/data/lerobot +export XTRAINER2_PICK_BOTTLE_PROVENANCE_PATH=$XTRAINER2_DATA_ROOT/xtrainer2_optr_support/pick_bottle/data_provenance.json +export XTRAINER2_PICK_BOTTLE_FINGERPRINT_PATH=$XTRAINER2_DATA_ROOT/xtrainer2_optr_support/pick_bottle/content_fingerprints.json +export OPTR_ACCELERATE_CONFIG=$OPTR_PROJECT_ROOT/configs/deepspeed_florence_8gpu.yaml +export OPTR_NUM_PROCESSES=8 + +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_RUN_ID="${WANDB_RUN_ID:-pick-bottle-rel-aug-60ep}" +export WANDB_RESUME=allow + +run_id=florence-xtrainer2-pick-bottle-rel-aug-uniform +log=$OPTR_RUN_ROOT/${run_id}.log +mkdir -p "$OPTR_RUN_ROOT" + +echo "launching ${run_id}" +echo " log: ${log}" +exec bash examples/Xtrainer2/train_files/run_xtrainer2_train.sh \ + configs/florence_xtrainer2_pick_bottle_rel_aug_uniform.yaml \ + "$run_id" >> "$log" 2>&1 diff --git a/real_robot/source/scripts/run_xtrainer2_arm.sh b/real_robot/source/scripts/run_xtrainer2_arm.sh new file mode 100644 index 0000000000000000000000000000000000000000..a4adf8e5a65357f2d71a76e0d6dd3a3c739f0849 --- /dev/null +++ b/real_robot/source/scripts/run_xtrainer2_arm.sh @@ -0,0 +1,169 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -lt 2 || $# -gt 4 ]]; then + echo "usage: $0 [proposal.json] [config.yaml]" >&2 + exit 64 +fi + +arm=$1 +run_id=$2 +proposal=${3:-${OPTR_TEMPORAL_PROPOSAL_PATH:-}} +project="${OPTR_PROJECT_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" +python="${OPTR_PYTHON:-python}" +config=${4:-"${project}/configs/florence_xtrainer2_pack_phone_residual_20ep_to_60ep.yaml"} +accelerate_config="${OPTR_ACCELERATE_CONFIG:-${project}/configs/deepspeed_florence_8gpu.yaml}" +num_processes="${OPTR_NUM_PROCESSES:-8}" + +missing=() +for var in OPTR_RUN_ROOT FLORENCE_MODEL_PATH XTRAINER2_DATA_ROOT \ + XTRAINER2_PACK_PHONE_PROVENANCE_PATH XTRAINER2_PACK_PHONE_FINGERPRINT_PATH \ + XTRAINER2_PACK_PHONE_MANIFEST_PATH XTRAINER2_PACK_PHONE_RESUME_DIR \ + XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH; do + [[ -n "${!var:-}" ]] || missing+=("$var") +done +if (( ${#missing[@]} )); then + echo "error: required environment variables are unset: ${missing[*]}" >&2 + exit 78 +fi + +[[ "$num_processes" == 8 ]] || { + echo "error: OPTR continuation requires exactly 8 processes (got $num_processes)" >&2 + exit 78 +} +[[ "${FLORENCE_MODEL_PATH,,}" == *florence* && -d "$FLORENCE_MODEL_PATH" ]] || { + echo "error: FLORENCE_MODEL_PATH must be an existing Florence directory" >&2 + exit 66 +} +[[ -d "$XTRAINER2_DATA_ROOT" ]] || { + echo "error: XTRAINER2_DATA_ROOT is not a directory: $XTRAINER2_DATA_ROOT" >&2 + exit 66 +} +for path in "$config" "$accelerate_config" \ + "$XTRAINER2_PACK_PHONE_PROVENANCE_PATH" \ + "$XTRAINER2_PACK_PHONE_FINGERPRINT_PATH" \ + "$XTRAINER2_PACK_PHONE_MANIFEST_PATH" \ + "$XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH"; do + [[ -f "$path" ]] || { echo "error: required file missing: $path" >&2; exit 66; } +done +[[ -d "$XTRAINER2_PACK_PHONE_RESUME_DIR" ]] || { + echo "error: resume checkpoint directory missing: $XTRAINER2_PACK_PHONE_RESUME_DIR" >&2 + exit 66 +} +[[ "$(basename "$XTRAINER2_PACK_PHONE_RESUME_DIR")" == steps_17460_accelerate ]] || { + echo "error: exact common-start must be steps_17460_accelerate" >&2 + exit 78 +} +if ! command -v "$python" >/dev/null 2>&1 && [[ ! -x "$python" ]]; then + echo "error: python interpreter not found: $python" >&2 + exit 69 +fi + +case "$arm" in + uniform) + [[ -z "$proposal" ]] || { + echo "error: uniform arm does not accept a proposal" >&2 + exit 64 + } + proposal_args=(--datasets.vla_data.temporal_proposal_path=null) + ;; + shuffle|residual) + [[ -n "$proposal" && -f "$proposal" ]] || { + echo "error: $arm requires a proposal file via CLI or OPTR_TEMPORAL_PROPOSAL_PATH" >&2 + exit 66 + } + proposal_args=(--datasets.vla_data.temporal_proposal_path="$proposal") + ;; + *) + echo "error: unknown arm: $arm" >&2 + exit 64 + ;; +esac + +output_dir="${OPTR_RUN_ROOT}/${run_id}" +[[ ! -e "$output_dir" ]] || { + echo "error: refusing to reuse run directory: $output_dir" >&2 + exit 73 +} + +export PYTHONPATH="${project}:${PYTHONPATH:-}" +"$python" - "$XTRAINER2_PACK_PHONE_RESUME_DIR" \ + "$XTRAINER2_PACK_PHONE_RESUME_INTEGRITY_PATH" \ + "$XTRAINER2_PACK_PHONE_MANIFEST_PATH" "${proposal:-}" "$arm" <<'PY' +import json +import sys +from pathlib import Path + +checkpoint, inventory_path, manifest_path, proposal_path = map(Path, sys.argv[1:5]) +arm = sys.argv[5] +commit_path = checkpoint / "checkpoint.json" +try: + commit = json.loads(commit_path.read_text(encoding="utf-8")) +except (OSError, json.JSONDecodeError) as exc: + raise SystemExit(f"error: invalid checkpoint commit {commit_path}: {exc}") +if commit.get("schema_version") != 1: + raise SystemExit("error: unsupported checkpoint schema") +if int(commit.get("completed_steps", -1)) != 17460: + raise SystemExit("error: checkpoint is not global step 17460") +if int(commit.get("world_size", -1)) != 8: + raise SystemExit("error: checkpoint world size is not 8") +files = commit.get("files") +if not isinstance(files, list) or not files: + raise SystemExit("error: checkpoint commit has no file inventory") +for entry in files: + path = checkpoint / entry["name"] + if not path.is_file() or path.stat().st_size != int(entry["bytes"]): + raise SystemExit(f"error: incomplete checkpoint file: {path}") + +try: + inventory = json.loads(inventory_path.read_text(encoding="utf-8")) +except (OSError, json.JSONDecodeError) as exc: + raise SystemExit(f"error: invalid checkpoint integrity manifest: {exc}") +expected = [(x["name"], int(x["bytes"])) for x in files] +if inventory.get("schema_version") != 1: + raise SystemExit("error: checkpoint integrity manifest must use schema 1") +if inventory.get("status") == "locked": + observed = [(x["name"], int(x["bytes"])) for x in inventory.get("files", [])] + if observed != expected: + raise SystemExit("error: checkpoint integrity manifest differs from checkpoint.json") +elif inventory.get("status") == "pass": + if ( + inventory.get("errors") != [] + or int(inventory.get("completed_steps", -1)) != 17460 + or int(inventory.get("world_size", -1)) != 8 + or int(inventory.get("inventory_files", -1)) != len(files) + ): + raise SystemExit("error: common-start preflight report is inconsistent") +else: + raise SystemExit("error: checkpoint integrity manifest did not pass") + +try: + from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest + from starVLA.dataloader.temporal_proposal import TemporalProposal + manifest = ValidAnchorManifest.load(manifest_path) + if arm != "uniform": + proposal = TemporalProposal.load(proposal_path, manifest) + if proposal.score_kind != arm: + raise ValueError( + f"{arm} arm received {proposal.score_kind!r} proposal" + ) +except Exception as exc: + raise SystemExit(f"error: manifest/proposal validation failed: {exc}") +PY + +if [[ "${OPTR_PREFLIGHT_ONLY:-0}" == 1 ]]; then + echo "preflight passed: $arm -> $output_dir" + exit 0 +fi + +cd "$project" +export WANDB_MODE="${WANDB_MODE:-offline}" +export PYTHONUNBUFFERED=1 +exec "$python" -m accelerate.commands.launch \ + --config_file "$accelerate_config" \ + --num_processes 8 \ + starVLA/training/train_starvla.py \ + --config_yaml "$config" \ + --run_id "$run_id" \ + --run_root_dir "$OPTR_RUN_ROOT" \ + "${proposal_args[@]}" diff --git a/real_robot/source/scripts/validate_xtrainer2_lerobot.py b/real_robot/source/scripts/validate_xtrainer2_lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..112afc5fe98205c444e69f5b33fd4d9162914f44 --- /dev/null +++ b/real_robot/source/scripts/validate_xtrainer2_lerobot.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Validate a converted xtrainer2 LeRobot dataset. + +Two independent checks: + +1. Decoded rotations must reproduce the official ``norm_stats.json`` in the + 20-D model space, which is what proves the quaternion convention is right. +2. StarVLA's dataloader must load a sample end to end. +""" + +from __future__ import annotations + +import argparse +import glob +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset +from starVLA.dataloader.gr00t_lerobot.transform.state_action import RotationTransform + +REPO_ROOT = Path(__file__).resolve().parents[1] +OFFICIAL_STATS = "/apdcephfs_nj7/share_305204761/arianliu/data/benchmark/norm_stats.json" +ARM_BLOCKS = ((0, 3, 7), (8, 11, 15)) + + +def load_data_config(): + """The examples tree is not a package, so load the registry by path.""" + path = REPO_ROOT / "examples/Xtrainer2/train_files/data_registry/data_config.py" + spec = importlib.util.spec_from_file_location("_xtrainer2_data_config", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module.Xtrainer2BimanualEefDataConfig() + + +def to_model_space(raw: np.ndarray, scramble_back_to_xyzw: bool = False) -> np.ndarray: + """Expand stored 16-D poses into the 20-D layout used by the official stats.""" + transform = RotationTransform("quaternion", "rotation_6d") + columns = [] + for position, rotation, gripper in ARM_BLOCKS: + columns.append(raw[:, position : position + 3]) + block = raw[:, rotation : rotation + 4] + if scramble_back_to_xyzw: + block = block[:, [1, 2, 3, 0]] + six = transform.forward(torch.from_numpy(np.ascontiguousarray(block))) + columns.append(six.numpy()) + columns.append(raw[:, gripper : gripper + 1]) + return np.concatenate(columns, axis=1) + + +def _stats_error(model_space: np.ndarray, official: dict) -> tuple[float, np.ndarray]: + ref_mean = np.array(official["state"]["mean"])[0] + ref_lo = np.array(official["state"]["p_low"])[0] + ref_hi = np.array(official["state"]["p_high"])[0] + mean = model_space.mean(0) + lo = np.quantile(model_space, 0.01, axis=0) + hi = np.quantile(model_space, 0.99, axis=0) + error = float( + (np.abs(mean - ref_mean).mean() + np.abs(lo - ref_lo).mean() + np.abs(hi - ref_hi).mean()) + / 3 + ) + return error, mean + + +def read_states(dataset_root: Path) -> np.ndarray: + frames = [ + pd.read_parquet(path)["observation.state"] + for path in sorted(glob.glob(str(dataset_root / "data" / "chunk-*" / "*.parquet"))) + ] + return np.stack([np.asarray(row, dtype=np.float32) for f in frames for row in f]) + + +def check_convention(dataset_roots: list[Path]) -> tuple[float, float]: + """Compare the stored convention against the wrong one, pooled over all tasks. + + The official stats cover the whole benchmark, so any single task is biased + against them; only the pooled corpus is comparable, and only the gap between + the two conventions is meaningful. + """ + official = json.load(open(OFFICIAL_STATS)) + ref_mean = np.array(official["state"]["mean"])[0] + raw = np.concatenate([read_states(root) for root in dataset_roots], axis=0) + + stored_error, mean = _stats_error(to_model_space(raw), official) + wrong_error, _ = _stats_error(to_model_space(raw, scramble_back_to_xyzw=True), official) + + print(f" pooled frames={len(raw)} over {len(dataset_roots)} datasets") + print(f" ours mean {np.round(mean, 3)}") + print(f" ref mean {np.round(ref_mean, 3)}") + print(f" err as stored (wxyz) : {stored_error:.4f}") + print(f" err if misread (xyzw) : {wrong_error:.4f}") + return stored_error, wrong_error + + +def check_dataloader(dataset_root: Path) -> None: + config = load_data_config() + dataset = LeRobotSingleDataset( + dataset_path=dataset_root, + modality_configs=config.modality_config(), + embodiment_tag=config.embodiment_tag, + transforms=config.transform(), + data_cfg={"lerobot_version": "v2.0", "action_mode": "abs", "include_state": True}, + ) + sample = dataset[0] + print(f" episodes={len(dataset.trajectory_lengths)} samples={len(dataset)}") + for key, value in sample.items(): + if isinstance(value, (np.ndarray, torch.Tensor)): + print(f" {key}: {tuple(value.shape)} {value.dtype}") + else: + print(f" {key}: {value!r}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + type=Path, + default=Path("/apdcephfs_nj7/share_305204761/arianliu/data/lerobot"), + ) + parser.add_argument( + "--min-margin", + type=float, + default=3.0, + help="How many times better the stored convention must fit the official stats", + ) + args = parser.parse_args() + + dataset_roots = sorted(p for p in args.root.iterdir() if p.is_dir()) + failures = [] + + print("=== rotation convention vs official norm_stats ===") + stored_error, wrong_error = check_convention(dataset_roots) + margin = wrong_error / max(stored_error, 1e-9) + print(f" margin: {margin:.1f}x") + if margin < args.min_margin: + failures.append(f"convention margin {margin:.1f}x < {args.min_margin}x") + + for dataset_root in dataset_roots: + print(f"\n=== {dataset_root.name}: StarVLA dataloader ===") + check_dataloader(dataset_root) + + if failures: + print("\nFAILED:") + for failure in failures: + print(f" {failure}") + return 1 + print("\nAll checks passed.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/real_robot/source/starVLA/__init__.py b/real_robot/source/starVLA/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/real_robot/source/starVLA/dataloader/__init__.py b/real_robot/source/starVLA/dataloader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..58b002c639fa544d4b82b7f5e0d03c3a1b7dd669 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/__init__.py @@ -0,0 +1,104 @@ +# Modified in the ReTime research artifact; see NOTICE. +import json +import os +from accelerate.logging import get_logger +import numpy as np +from torch.utils.data import DataLoader +import torch.distributed as dist +from pathlib import Path +from starVLA.dataloader.infoflow_sampler import StatelessDrawSampler + +logger = get_logger(__name__) + +def save_dataset_statistics(dataset_statistics, run_dir): + """Saves a `dataset_statistics.json` file.""" + out_path = run_dir / "dataset_statistics.json" + with open(out_path, "w") as f_json: + for _, stats in dataset_statistics.items(): + for k in stats["action"].keys(): + if isinstance(stats["action"][k], np.ndarray): + stats["action"][k] = stats["action"][k].tolist() + if "proprio" in stats: + for k in stats["proprio"].keys(): + if isinstance(stats["proprio"][k], np.ndarray): + stats["proprio"][k] = stats["proprio"][k].tolist() + if "num_trajectories" in stats: + if isinstance(stats["num_trajectories"], np.ndarray): + stats["num_trajectories"] = stats["num_trajectories"].item() + if "num_transitions" in stats: + if isinstance(stats["num_transitions"], np.ndarray): + stats["num_transitions"] = stats["num_transitions"].item() + json.dump(dataset_statistics, f_json, indent=2) + logger.info(f"Saved dataset statistics file at path {out_path}") + + + +def build_dataloader(cfg, dataset_py="lerobot_datasets_oxe"): + + if dataset_py == "lerobot_datasets": + from starVLA.dataloader.lerobot_datasets import get_vla_dataset, collate_fn + vla_dataset_cfg = cfg.datasets.vla_data + + vla_dataset = get_vla_dataset( + data_cfg=vla_dataset_cfg, + balance_dataset_weights=vla_dataset_cfg.get("balance_dataset_weights", False), + balance_trajectory_weights=vla_dataset_cfg.get("balance_trajectory_weights", False), + ) + + traceable_sampling = bool(vla_dataset_cfg.get("traceable_sampling", False)) + per_device_batch_size = int(cfg.datasets.vla_data.per_device_batch_size) + world_size = ( + dist.get_world_size() + if dist.is_available() and dist.is_initialized() + else int(vla_dataset_cfg.get("expected_world_size", 1)) + ) + sampler = None + if traceable_sampling: + expected_world_size = int(vla_dataset_cfg.expected_world_size) + if world_size != expected_world_size: + raise RuntimeError( + f"Traceable sampling requires world size {expected_world_size}, " + f"but distributed runtime reports {world_size}" + ) + gradient_accumulation_steps = int( + cfg.trainer.gradient_accumulation_steps + ) + global_update_batch = ( + per_device_batch_size + * world_size + * gradient_accumulation_steps + ) + usable_draws = len(vla_dataset) - ( + len(vla_dataset) % global_update_batch + ) + if usable_draws <= 0: + raise ValueError( + f"Dataset length {len(vla_dataset)} is smaller than one global " + f"optimizer update ({global_update_batch} draws)" + ) + sampler = StatelessDrawSampler(usable_draws) + num_workers = int(vla_dataset_cfg.get("num_workers", 16)) + dataloader_kwargs = dict( + dataset=vla_dataset, + batch_size=per_device_batch_size, + collate_fn=collate_fn, + sampler=sampler, + num_workers=num_workers, + pin_memory=True, + persistent_workers=num_workers > 0, + drop_last=traceable_sampling, + ) + if num_workers > 0: + dataloader_kwargs["prefetch_factor"] = int(vla_dataset_cfg.get("prefetch_factor", 4)) + vla_train_dataloader = DataLoader(**dataloader_kwargs) + + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + if rank == 0: + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + vla_dataset.save_dataset_statistics(output_dir / "dataset_statistics.json") + if vla_dataset.valid_anchor_manifest is not None: + manifest_path = output_dir / "valid_anchor_manifest.json" + vla_dataset.valid_anchor_manifest.save(manifest_path) + return vla_train_dataloader + raise ValueError(f"Unsupported dataset loader: {dataset_py}") diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/__init__.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/data_config.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/data_config.py new file mode 100644 index 0000000000000000000000000000000000000000..12f13ce48d87a820eac02cc9c0ee7fe941234818 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/data_config.py @@ -0,0 +1,1100 @@ +# Modified in the ReTime research artifact; see NOTICE. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + + +from abc import ABC, abstractmethod + +from starVLA.dataloader.gr00t_lerobot.datasets import ModalityConfig +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag +from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform, ModalityTransform +from starVLA.dataloader.gr00t_lerobot.transform.concat import ConcatTransform +from starVLA.dataloader.gr00t_lerobot.transform.state_action import ( + StateActionSinCosTransform, + StateActionToTensor, + StateActionTransform, +) +from starVLA.dataloader.gr00t_lerobot.transform.video import ( + VideoColorJitter, + VideoCrop, + VideoResize, + VideoToNumpy, + VideoToTensor, +) +# from gr00t.model.transforms import GR00TTransform + + +class BaseDataConfig(ABC): + @abstractmethod + def modality_config(self) -> dict[str, ModalityConfig]: + pass + + @abstractmethod + def transform(self) -> ModalityTransform: + pass + + +########################################################################################### + +class OxeDroidDataConfig: + embodiment_tag = EmbodimentTag.OXE_DROID + video_keys = [ + "video.exterior_image_1", + "video.exterior_image_2", + "video.wrist_image", + ] + state_keys = [ + "state.eef_position", + "state.eef_rotation", + "state.gripper_position", + ] + action_keys = [ + "action.eef_position_delta", + "action.eef_rotation_delta", + "action.gripper_position", + ] + language_keys = ["annotation.language.language_instruction"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # video transforms + VideoToTensor(apply_to=self.video_keys), + VideoCrop(apply_to=self.video_keys, scale=0.95), + VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"), + VideoColorJitter( + apply_to=self.video_keys, + brightness=0.3, + contrast=0.4, + saturation=0.5, + hue=0.08, + ), + VideoToNumpy(apply_to=self.video_keys), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.eef_position": "min_max", + "state.gripper_position": "min_max", + }, + target_rotations={ + "state.eef_rotation": "rotation_6d", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.gripper_position": "binary", + }, + target_rotations={"action.eef_rotation_delta": "axis_angle"}, + ), + # concat transforms + ConcatTransform( + video_concat_order=self.video_keys, + state_concat_order=self.state_keys, + action_concat_order=self.action_keys, + ), + GR00TTransform( + state_horizon=len(self.observation_indices), + action_horizon=len(self.action_indices), + max_state_dim=64, + max_action_dim=32, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + + +class OxeBridgeDataConfig: + embodiment_tag = EmbodimentTag.OXE_BRIDGE + video_keys = [ + "video.image_0", + ] + state_keys = [ + "state.x", + "state.y", + "state.z", + "state.roll", + "state.pitch", + "state.yaw", + "state.pad", + "state.gripper", + ] + action_keys = [ + "action.x", + "action.y", + "action.z", + "action.roll", + "action.pitch", + "action.yaw", + "action.gripper", + ] + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # video transforms + # VideoToTensor(apply_to=self.video_keys), + # VideoCrop(apply_to=self.video_keys, scale=0.95), + # VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"), + # VideoColorJitter( + # apply_to=self.video_keys, + # brightness=0.3, + # contrast=0.4, + # saturation=0.5, + # hue=0.08, + # ), + # VideoToNumpy(apply_to=self.video_keys), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.x": "q99", + "state.y": "q99", + "state.z": "q99", + "state.roll": "q99", + "state.pitch": "q99", + "state.yaw": "q99", + "state.pad": "q99", + "state.gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "q99", + "action.y": "q99", + "action.z": "q99", + "action.roll": "q99", + "action.pitch": "q99", + "action.yaw": "q99", + "action.gripper": "binary", + }, + ), + # concat transforms + # ConcatTransform( + # # video_concat_order=self.video_keys, + # state_concat_order=self.state_keys, + # action_concat_order=self.action_keys, + # ), + # GR00TTransform( + # state_horizon=len(self.observation_indices), + # action_horizon=len(self.action_indices), + # max_state_dim=64, + # max_action_dim=32, + # ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + +class OxeRT1DataConfig: + embodiment_tag = EmbodimentTag.OXE_RT1 + video_keys = [ + "video.image", + ] + state_keys = [ + "state.x", + "state.y", + "state.z", + "state.rx", + "state.ry", + "state.rz", + "state.rw", + "state.gripper", + ] + action_keys = [ + "action.x", + "action.y", + "action.z", + "action.roll", + "action.pitch", + "action.yaw", + "action.gripper", + ] + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # video transforms + # VideoToTensor(apply_to=self.video_keys), + # VideoCrop(apply_to=self.video_keys, scale=0.95), + # VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"), + # VideoColorJitter( + # apply_to=self.video_keys, + # brightness=0.3, + # contrast=0.4, + # saturation=0.5, + # hue=0.08, + # ), + # VideoToNumpy(apply_to=self.video_keys), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.x": "q99", + "state.y": "q99", + "state.z": "q99", + "state.rx": "q99", + "state.ry": "q99", + "state.rz": "q99", + "state.rw": "q99", + "state.gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "q99", + "action.y": "q99", + "action.z": "q99", + "action.roll": "q99", + "action.pitch": "q99", + "action.yaw": "q99", + "action.gripper": "binary", + }, + ), + # concat transforms + # ConcatTransform( + # # video_concat_order=self.video_keys, + # state_concat_order=self.state_keys, + # action_concat_order=self.action_keys, + # ), + # GR00TTransform( + # state_horizon=len(self.observation_indices), + # action_horizon=len(self.action_indices), + # max_state_dim=64, + # max_action_dim=32, + # ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + + +class SingleFrankaRobotiqDeltaEefDataConfig: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.base_view", + "video.ego_view", + ] + state_keys = [ + "state.eef_position", + "state.eef_rotation", + ] + action_keys = [ + "action.delta_eef_position", + "action.delta_eef_rotation", + "action.gripper_close", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.eef_position": "min_max", + "state.eef_rotation": "min_max", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.delta_eef_position": "min_max", + "action.delta_eef_rotation": "min_max", + "action.gripper_close": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + +class Libero4in1DataConfig: + embodiment_tag = EmbodimentTag.FRANKA + video_keys = [ + "video.primary_image", + ] + + state_keys = [ + "state.x", + "state.y", + "state.z", + "state.roll", + "state.pitch", + "state.yaw", + "state.pad", + "state.gripper", + ] + action_keys = [ + "action.x", + "action.y", + "action.z", + "action.roll", + "action.pitch", + "action.yaw", + "action.gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + + observation_indices = [0] + action_indices = list(range(8)) + state_indices = list(range(-16, 0)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.state_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "min_max", + "action.y": "min_max", + "action.z": "min_max", + "action.roll": "min_max", + "action.pitch": "min_max", + "action.yaw": "min_max", + # LIBERO's gripper is a native 0/1 control, not a continuous + # min-max dimension. This preserves 0 -> 0 and 1 -> 1 while + # producing the official statistics mask value `false`. + "action.gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + + +class SingleFrankaRobotiqDeltaJointsDataConfig: + embodiment_tag = EmbodimentTag.FRANKA + video_keys = [ + "video.base_view", + "video.ego_view", + ] + state_keys = [ + "state.joints", + ] + action_keys = [ + "action.delta_joints", + "action.gripper_close", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.joints": "min_max", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.delta_joints": "min_max", + "action.gripper_close": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + +class FourierGr1ArmsWaistDataConfig: + embodiment_tag = EmbodimentTag.GR1 + video_keys = ["video.ego_view"] + state_keys = [ + "state.left_arm", + "state.right_arm", + "state.left_hand", + "state.right_hand", + "state.waist", + ] + action_keys = [ + "action.left_arm", + "action.right_arm", + "action.left_hand", + "action.right_hand", + "action.waist", + ] + language_keys = ["annotation.human.coarse_action"] + observation_indices = [0] + action_indices = list(range(16)) + + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self) -> ModalityTransform: + transforms = [ + # video transforms + # VideoToTensor(apply_to=self.video_keys), + # VideoCrop(apply_to=self.video_keys, scale=0.95), + # VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"), + # VideoColorJitter( + # apply_to=self.video_keys, + # brightness=0.3, + # contrast=0.4, + # saturation=0.5, + # hue=0.08, + # ), + # VideoToNumpy(apply_to=self.video_keys), + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionSinCosTransform(apply_to=self.state_keys), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={key: "min_max" for key in self.action_keys}, + ), + # concat transforms + # ConcatTransform( + # video_concat_order=self.video_keys, + # state_concat_order=self.state_keys, + # action_concat_order=self.action_keys, + # ), + ] + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + + +########################################################################################### + +class SO101Config: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + #input + video_keys = [ + "video.primary_image", + "video.wrist_image", + ] + + state_keys = [ + "state.shoulder_pan.pos", + "state.shoulder_lift.pos", + "state.elbow_flex.pos", + "state.wrist_flex.pos", + "state.wrist_roll.pos", + "state.gripper.pos", + ] + language_keys = ["annotation.human.action.task_description"] + + # output + action_keys = [ + "action.shoulder_pan.pos", + "action.shoulder_lift.pos", + "action.elbow_flex.pos", + "action.wrist_flex.pos", + "action.wrist_roll.pos", + "action.gripper.pos", + ] + + + observation_indices = [0] + action_indices = list(range(16)) + + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + key: "min_max" for key in self.state_keys + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + key: "min_max" for key in self.action_keys + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + + + +class ArxX5DataConfig: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_joints", + "state.right_joints", + "state.left_gripper", + "state.right_gripper", + ] + action_keys = [ + "action.left_joints", + "action.right_joints", + "action.left_gripper", + "action.right_gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.left_joints": "min_max", + "state.right_joints": "min_max", + "state.left_gripper": "binary", + "state.right_gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.left_joints": "min_max", + "action.right_joints": "min_max", + "action.left_gripper": "binary", + "action.right_gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + + +class AgilexDataConfig: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_joints", + "state.right_joints", + "state.left_gripper", + "state.right_gripper", + ] + action_keys = [ + "action.left_joints", + "action.right_joints", # This order differs from the dataset schema. + "action.left_gripper", + "action.right_gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(16)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + normalization_modes={ + "state.left_joints": "min_max", + "state.right_joints": "min_max", + "state.left_gripper": "binary", + "state.right_gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.left_joints": "min_max", + "action.right_joints": "min_max", + "action.left_gripper": "binary", + "action.right_gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + + +class AgilexData50Config: + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + video_keys = [ + "video.cam_high", + "video.cam_left_wrist", + "video.cam_right_wrist", + ] + state_keys = [ + "state.left_joints", + "state.right_joints", + "state.left_gripper", + "state.right_gripper", + ] + action_keys = [ + "action.left_joints", + "action.right_joints", # This order differs from the dataset schema. + "action.left_gripper", + "action.right_gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + observation_indices = [0] + action_indices = list(range(50)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + modality_configs = { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + return modality_configs + + def transform(self): + transforms = [ + # state transforms + StateActionToTensor(apply_to=self.state_keys), + StateActionTransform( + apply_to=self.state_keys, + binary_threshold=0.49, + normalization_modes={ + "state.left_joints": "min_max", + "state.right_joints": "min_max", + "state.left_gripper": "binary", + "state.right_gripper": "binary", + }, + ), + # action transforms + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + binary_threshold=0.49, + normalization_modes={ + "action.left_joints": "min_max", + "action.right_joints": "min_max", + "action.left_gripper": "binary", + "action.right_gripper": "binary", + }, + ), + ] + + return ComposedModalityTransform(transforms=transforms) + +########################################################################################### + +class VLAArenaFrankaDataConfig: + """ + Data config for VLA-Arena environments (robosuite, mounted Panda). + + Action space : 7-DOF delta EEF (Δx, Δy, Δz, Δroll, Δpitch, Δyaw, gripper) + Observation : primary agentview image (256×256 rendered, resized to 224×224) + State : EEF pos (3) + EEF axis-angle (3) + gripper qpos (1) = 7 + """ + + embodiment_tag = EmbodimentTag.FRANKA + video_keys = [ + "video.primary_image", # agentview camera + ] + state_keys = [ + "state.x", + "state.y", + "state.z", + "state.roll", + "state.pitch", + "state.yaw", + "state.gripper", + ] + action_keys = [ + "action.x", + "action.y", + "action.z", + "action.roll", + "action.pitch", + "action.yaw", + "action.gripper", + ] + + language_keys = ["annotation.human.action.task_description"] + + observation_indices = [0] + action_indices = list(range(8)) + state_indices = list(range(-16, 0)) + + def modality_config(self): + video_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.video_keys, + ) + state_modality = ModalityConfig( + delta_indices=self.state_indices, + modality_keys=self.state_keys, + ) + action_modality = ModalityConfig( + delta_indices=self.action_indices, + modality_keys=self.action_keys, + ) + language_modality = ModalityConfig( + delta_indices=self.observation_indices, + modality_keys=self.language_keys, + ) + return { + "video": video_modality, + "state": state_modality, + "action": action_modality, + "language": language_modality, + } + + def transform(self): + transforms = [ + StateActionToTensor(apply_to=self.action_keys), + StateActionTransform( + apply_to=self.action_keys, + normalization_modes={ + "action.x": "min_max", + "action.y": "min_max", + "action.z": "min_max", + "action.roll": "min_max", + "action.pitch": "min_max", + "action.yaw": "min_max", + }, + ), + ] + return ComposedModalityTransform(transforms=transforms) + + +########################################################################################### + +ROBOT_TYPE_CONFIG_MAP = { + "libero_franka": Libero4in1DataConfig(), + "oxe_droid": OxeDroidDataConfig(), + "oxe_bridge": OxeBridgeDataConfig(), + "oxe_rt1": OxeRT1DataConfig(), + "SO101": SO101Config(), + "demo_sim_franka_delta_joints": SingleFrankaRobotiqDeltaJointsDataConfig(), + "arx_x5": ArxX5DataConfig(), + "robotwin": AgilexDataConfig(), + "robotwin50": AgilexData50Config(), + "fourier_gr1_arms_waist": FourierGr1ArmsWaistDataConfig(), + "vla_arena_franka": VLAArenaFrankaDataConfig(), + + "custom_robot_config": SingleFrankaRobotiqDeltaEefDataConfig(), +} diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/datasets.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..39b235b256b3924fc18c7786c5321f2ceb29ea8d --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/datasets.py @@ -0,0 +1,3280 @@ +# Modified in the ReTime research artifact; see NOTICE. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + + +""" +In this file, we define 3 types of datasets: +1. LeRobotSingleDataset: a single dataset for a given embodiment tag +2. LeRobotMixtureDataset: a mixture of datasets for a given list of embodiment tags +3. CachedLeRobotSingleDataset: a single dataset for a given embodiment tag, + with caching for the video frames + +See `scripts/load_dataset.py` for examples on how to use these datasets. +""" +import os +from contextlib import contextmanager +import hashlib +import io +import json, torch +import copy +from collections import defaultdict +from pathlib import Path +from typing import Sequence +import os, random +import numpy as np +import pandas as pd +from pydantic import BaseModel, Field, ValidationError +from torch.utils.data import Dataset +from tqdm import tqdm +from PIL import Image +import torch.distributed as dist + +from starVLA.dataloader.gr00t_lerobot.video import get_all_frames, get_frames_by_timestamps + +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag +from starVLA.dataloader.gr00t_lerobot.schema import ( + DatasetMetadata, + DatasetStatisticalValues, + LeRobotModalityMetadata, + LeRobotStateActionMetadata, +) +from starVLA.dataloader.gr00t_lerobot.transform import ComposedModalityTransform +from starVLA.dataloader.gr00t_lerobot.transform.state_action import StateActionTransform +from starVLA.dataloader.infoflow_sampler import ( + LogicalDraw, + ValidAnchorManifest, + compute_dataset_content_fingerprint, + validate_dataset_content_inventory, +) +from starVLA.dataloader.temporal_proposal import TemporalProposal +from starVLA.dataloader.relative_pose import relative_values, rotation_type_is_quaternion + +from functools import partial +from typing import Tuple, List +import pickle +import gc + +# LeRobot v2.0 dataset file names +LE_ROBOT_MODALITY_FILENAME = "meta/modality.json" +LE_ROBOT_EPISODE_FILENAME = "meta/episodes.jsonl" +LE_ROBOT_TASKS_FILENAME = "meta/tasks.jsonl" +LE_ROBOT_INFO_FILENAME = "meta/info.json" +LE_ROBOT_STATS_FILENAME = "meta/stats_gr00t.json" +LE_ROBOT_DATA_FILENAME = "data/*/*.parquet" +LE_ROBOT_STEPS_FILENAME = "meta/steps.pkl" +LE_ROBOT_STATS_FORMAT_VERSION = 2 +EPSILON = 5e-4 + + +@contextmanager +def _isolated_transform_rng(seed: int | None): + """Make data augmentation a pure function of one logical draw id.""" + + if seed is None: + yield + return + python_state = random.getstate() + numpy_state = np.random.get_state() + torch_state = torch.get_rng_state() + bounded_seed = int(seed) % (2**32) + random.seed(bounded_seed) + np.random.seed(bounded_seed) + torch.manual_seed(bounded_seed) + try: + yield + finally: + random.setstate(python_state) + np.random.set_state(numpy_state) + torch.set_rng_state(torch_state) + +# LeRobot v3.0 dataset file names +LE_ROBOT3_TASKS_FILENAME = "meta/tasks.parquet" +LE_ROBOT3_EPISODE_FILENAME = "meta/episodes/*/*.parquet" + + +def calculate_dataset_statistics(parquet_paths: list[Path]) -> dict: + """Calculate the dataset statistics of all columns for a list of parquet files.""" + # Dataset statistics + all_low_dim_data_list = [] + # Collect all the data + # parquet_paths = parquet_paths[:3] + for parquet_path in tqdm( + sorted(list(parquet_paths)), + desc="Collecting all parquet files...", + ): + # Load the parquet file + parquet_data = pd.read_parquet(parquet_path) + parquet_data = parquet_data + all_low_dim_data_list.append(parquet_data) + + if not all_low_dim_data_list: + raise FileNotFoundError( + f"No parquet files found under the provided paths: {[str(p) for p in parquet_paths[:3]]}..." + f" — make sure the dataset has been downloaded/converted before training." + ) + + all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) + # Compute dataset statistics + dataset_statistics = {} + for le_modality in tqdm(all_low_dim_data.columns, desc="Processing modalities"): + print(le_modality) + if "task_info" in le_modality: + continue + print(f"Computing statistics for {le_modality}...") + try: + np_data = np.vstack( + [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]] + ) + except Exception as e: + print(f"Warning: Failed to process modality {le_modality} due to error: {e}") + continue + + dataset_statistics[le_modality] = { + "mean": np.mean(np_data, axis=0).tolist(), + "std": np.std(np_data, axis=0).tolist(), + "min": np.min(np_data, axis=0).tolist(), + "max": np.max(np_data, axis=0).tolist(), + "q01": np.quantile(np_data, 0.01, axis=0).tolist(), + "q99": np.quantile(np_data, 0.99, axis=0).tolist(), + } + return dataset_statistics + + +def _normalize_action_mode(mode: str) -> str: + """Normalize action mode names to {abs, delta, rel}.""" + # Keep action-mode logic isolated from the generic dataset path. + mode = str(mode).lower() + if mode in {"absolute", "raw"}: + mode = "abs" + if mode not in {"abs", "delta", "rel"}: + mode = "abs" + return mode + + +def _normalize_action_mode_apply_keys( + action_mode_apply_keys: Sequence[str] | None, + fallback_keys: Sequence[str] | None = None, +) -> list[str]: + source_keys = action_mode_apply_keys if action_mode_apply_keys else (fallback_keys or []) + normalized = [] + for key in source_keys: + key = str(key) + if not key.startswith("action."): + key = f"action.{key}" + normalized.append(key) + return normalized + + +def _normalize_action_mode_state_map(action_mode_state_map: dict[str, str] | None) -> dict[str, str]: + normalized = {} + for action_key, state_key in (action_mode_state_map or {}).items(): + action_key = str(action_key) + state_key = str(state_key) + if not action_key.startswith("action."): + action_key = f"action.{action_key}" + if not state_key.startswith("state."): + state_key = f"state.{state_key}" + normalized[action_key] = state_key + return normalized + + +def _build_stats_cache_config( + action_mode: str, +) -> dict: + return { + "mode": action_mode, + } + + +def _invalidate_legacy_stats_cache(stats_path: Path, reason: str) -> None: + if not stats_path.exists(): + return + print(f"Removing stale dataset statistics cache at {stats_path}: {reason}") + stats_path.unlink() + + +def _load_stats_cache( + stats_path: Path, + expected_config: dict, + *, + invalidate_legacy: bool, +) -> dict | None: + if not stats_path.exists(): + return None + + try: + with open(stats_path, "r") as f: + payload = json.load(f) + except Exception as exc: + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, f"failed to load JSON ({exc})") + return None + + if not isinstance(payload, dict): + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, "unexpected top-level format") + return None + + format_version = payload.get("__format_version") + cache_config = payload.get("__cache_config") + statistics = payload.get("statistics") + if format_version != LE_ROBOT_STATS_FORMAT_VERSION or cache_config is None or statistics is None: + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, "legacy statistics format detected") + return None + + if cache_config != expected_config: + if invalidate_legacy: + _invalidate_legacy_stats_cache(stats_path, "statistics config mismatch, rebuilding cache") + return None + + return statistics + + +def _save_stats_cache(stats_path: Path, cache_config: dict, statistics: dict) -> None: + payload = { + "__format_version": LE_ROBOT_STATS_FORMAT_VERSION, + "__cache_config": cache_config, + "statistics": statistics, + } + tmp_path = stats_path.with_suffix(".tmp") + stats_path.parent.mkdir(parents=True, exist_ok=True) + with open(tmp_path, "w") as f: + json.dump(payload, f, indent=4) + os.replace(tmp_path, stats_path) + + +def _compute_statistics_for_mode( + parquet_paths: list[Path], + dataset_name: str, + action_mode: str, + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int] | None, + state_indices: list[int] | None, + action_mode_apply_keys: list[str] | None, + action_mode_state_map: dict[str, str] | None, +) -> dict: + if int(os.environ.get("RANK", "0")) == 0: + print(f"[RANK 0] Calculating dataset statistics for {dataset_name} (mode={action_mode})") + + base_stats = calculate_dataset_statistics(parquet_paths) + + if action_mode == "abs": + return base_stats + + if action_indices is None or state_indices is None: + raise ValueError( + "Both action and state modalities are required to compute " + f"{action_mode} action mode statistics." + ) + + if action_mode == "delta": + return calculate_delta_action_statistics( + parquet_paths=parquet_paths, + lerobot_modality_meta=lerobot_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=action_mode_apply_keys, + action_mode_state_map=action_mode_state_map, + base_stats=base_stats, + ) + if action_mode == "rel": + return calculate_rel_action_statistics( + parquet_paths=parquet_paths, + lerobot_modality_meta=lerobot_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=action_mode_apply_keys, + action_mode_state_map=action_mode_state_map, + base_stats=base_stats, + ) + raise ValueError(f"Unsupported action mode for statistics: {action_mode}") + + +def _load_or_compute_statistics( + stats_path: Path, + stats_cache_config: dict, + parquet_paths: list[Path], + dataset_name: str, + action_mode: str, + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int] | None, + state_indices: list[int] | None, + action_mode_apply_keys: list[str] | None, + action_mode_state_map: dict[str, str] | None, +) -> dict: + le_statistics = _load_stats_cache( + stats_path, + stats_cache_config, + invalidate_legacy=True, + ) + if le_statistics is not None: + return le_statistics + + le_statistics = _compute_statistics_for_mode( + parquet_paths=parquet_paths, + dataset_name=dataset_name, + action_mode=action_mode, + lerobot_modality_meta=lerobot_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=action_mode_apply_keys, + action_mode_state_map=action_mode_state_map, + ) + _save_stats_cache(stats_path, stats_cache_config, le_statistics) + return le_statistics + + +def _get_action_col_slices( + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_mode_apply_keys: list[str] | None = None, + action_mode_state_map: dict[str, str] | None = None, +) -> dict[str, list[tuple[tuple[int, int], str, tuple[int, int], str, str]]]: + apply_keys = _normalize_action_mode_apply_keys(action_mode_apply_keys, action_keys_full) + action_mode_state_map = _normalize_action_mode_state_map(action_mode_state_map) + + action_meta = lerobot_modality_meta.action + state_meta = lerobot_modality_meta.state + + # Build per-column mapping: action column -> list of (action_slice, state_column, state_slice) + action_col_slices: dict[str, list[tuple[tuple[int, int], str, tuple[int, int]]]] = {} + for action_key in apply_keys: + if not action_key.startswith("action."): + raise ValueError(f"Invalid action key {action_key}. Expected prefix 'action.'.") + state_key = action_mode_state_map.get(action_key, action_key.replace("action.", "state.", 1)) + if state_key not in state_keys_full: + raise ValueError( + f"State key {state_key} not found for action key {action_key}. " + f"Add it to action_mode_state_map or remove {action_key} from action_mode_apply_keys." + ) + + action_subkey = action_key.replace("action.", "", 1) + state_subkey = state_key.replace("state.", "", 1) + if action_subkey not in action_meta or state_subkey not in state_meta: + raise ValueError(f"Action/state key missing in metadata: {action_key} -> {state_key}") + + action_cfg = action_meta[action_subkey] + state_cfg = state_meta[state_subkey] + action_col = action_cfg.original_key or action_subkey + state_col = state_cfg.original_key or state_subkey + action_slice = (action_cfg.start, action_cfg.end) + state_slice = (state_cfg.start, state_cfg.end) + action_padding = "first_last" if action_cfg.absolute else "zero" + state_padding = "first_last" if state_cfg.absolute else "zero" + action_col_slices.setdefault(action_col, []).append( + (action_slice, state_col, state_slice, action_padding, state_padding) + ) + + return action_col_slices + + +def _action_slice_is_quaternion( + lerobot_modality_meta: "LeRobotModalityMetadata", + action_slice: tuple[int, int], +) -> bool: + start, end = action_slice + for cfg in lerobot_modality_meta.action.values(): + if cfg.start == start and cfg.end == end: + return rotation_type_is_quaternion(cfg.rotation_type) + return False + + +def calculate_delta_action_statistics( + parquet_paths: list[Path], + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int], + state_indices: list[int], + action_mode_apply_keys: list[str] | None = None, + action_mode_state_map: dict[str, str] | None = None, + base_stats: dict | None = None, +) -> dict: + """ + Calculate action statistics using delta mode. + + Rule: + - For t>0: a_t - a_{t-1} + - For t=0: a_0 - s_0 + + Mapping rule (only two cases): + 1) Use explicit action_mode_state_map if provided. + 2) Otherwise, replace 'action.' with 'state.' directly. + """ + if base_stats is None: + base_stats = calculate_dataset_statistics(parquet_paths) + + action_col_slices = _get_action_col_slices( + lerobot_modality_meta, action_keys_full, state_keys_full, action_mode_apply_keys, action_mode_state_map + ) + if not action_col_slices: + raise ValueError("No action columns found in the dataset.") + + def _get_chunk(array: np.ndarray, step_indices: np.ndarray, padding_strategy: str) -> np.ndarray: + max_length = array.shape[0] + front_padding = step_indices < 0 + end_padding = step_indices >= max_length + padding_positions = np.logical_or(front_padding, end_padding) + output = np.zeros((len(step_indices), array.shape[1]), dtype=array.dtype) + if (~padding_positions).any(): + output[~padding_positions] = array[step_indices[~padding_positions]] + if padding_positions.any(): + if padding_strategy == "first_last": + output[front_padding] = array[0] + output[end_padding] = array[-1] + elif padding_strategy == "zero": + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + accum: dict[str, list[np.ndarray]] = {col: [] for col in action_col_slices.keys()} + for parquet_path in tqdm(sorted(list(parquet_paths)), desc="Collecting delta action stats"): + data = pd.read_parquet(parquet_path) + trajectory_length = len(data) + for action_col, slice_list in action_col_slices.items(): + if action_col not in data.columns: + raise ValueError(f"{action_col} not found in parquet columns.") + action_matrix = np.stack(data[action_col]) + action_padding_ref = slice_list[0][3] + prepared_slices = [] + for a_slice, state_col, s_slice, action_padding, state_padding in slice_list: + if state_col not in data.columns: + raise ValueError(f"{state_col} not found in parquet columns.") + state_matrix = np.stack(data[state_col]) + state_part_full = state_matrix[:, s_slice[0] : s_slice[1]] + prepared_slices.append((a_slice, state_part_full, state_padding)) + for base_index in range(trajectory_length): + action_steps = np.array(action_indices) + base_index + action_chunk_full = _get_chunk(action_matrix, action_steps, action_padding_ref) + + for a_slice, state_part_full, state_padding in prepared_slices: + action_part_chunk = action_chunk_full[:, a_slice[0] : a_slice[1]] + state_chunk = _get_chunk(state_part_full, np.array(state_indices) + base_index, state_padding) + if action_part_chunk.shape[1] != state_chunk.shape[1]: + raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}") + + out = action_part_chunk.copy() + if len(out) > 1: + out[1:] = action_part_chunk[1:] - action_part_chunk[:-1] + out[0] = action_part_chunk[0] - state_chunk[0] + action_chunk_full[:, a_slice[0] : a_slice[1]] = out + + accum[action_col].append(action_chunk_full) + + delta_stats = copy.deepcopy(base_stats) + for action_col, series_list in accum.items(): + if not series_list: + continue + all_values = np.concatenate(series_list, axis=0).astype(np.float32) + delta_stats[action_col] = { + "mean": np.mean(all_values, axis=0).tolist(), + "std": np.std(all_values, axis=0).tolist(), + "min": np.min(all_values, axis=0).tolist(), + "max": np.max(all_values, axis=0).tolist(), + "q01": np.quantile(all_values, 0.01, axis=0).tolist(), + "q99": np.quantile(all_values, 0.99, axis=0).tolist(), + } + return delta_stats + + +def calculate_rel_action_statistics( + parquet_paths: list[Path], + lerobot_modality_meta: "LeRobotModalityMetadata", + action_keys_full: list[str], + state_keys_full: list[str], + action_indices: list[int], + state_indices: list[int], + action_mode_apply_keys: list[str] | None = None, + action_mode_state_map: dict[str, str] | None = None, + base_stats: dict | None = None, +) -> dict: + """ + Calculate action statistics using rel mode. + + Rule: + - Positions / scalars: a_t - s_0 + - Quaternion keys: q_rel = inv(q_s0) ⊗ q_a_t (still a unit quaternion) + + Mapping rule (only two cases): + 1) Use explicit action_mode_state_map if provided. + 2) Otherwise, replace 'action.' with 'state.' directly. + """ + if base_stats is None: + base_stats = calculate_dataset_statistics(parquet_paths) + + action_col_slices = _get_action_col_slices( + lerobot_modality_meta, action_keys_full, state_keys_full, action_mode_apply_keys, action_mode_state_map + ) + if not action_col_slices: + raise ValueError("No action columns found in the dataset.") + + def _get_chunk(array: np.ndarray, step_indices: np.ndarray, padding_strategy: str) -> np.ndarray: + max_length = array.shape[0] + front_padding = step_indices < 0 + end_padding = step_indices >= max_length + padding_positions = np.logical_or(front_padding, end_padding) + output = np.zeros((len(step_indices), array.shape[1]), dtype=array.dtype) + if (~padding_positions).any(): + output[~padding_positions] = array[step_indices[~padding_positions]] + if padding_positions.any(): + if padding_strategy == "first_last": + output[front_padding] = array[0] + output[end_padding] = array[-1] + elif padding_strategy == "zero": + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + accum: dict[str, list[np.ndarray]] = {col: [] for col in action_col_slices.keys()} + for parquet_path in tqdm(sorted(list(parquet_paths)), desc="Collecting rel action stats"): + data = pd.read_parquet(parquet_path) + trajectory_length = len(data) + for action_col, slice_list in action_col_slices.items(): + if action_col not in data.columns: + raise ValueError(f"{action_col} not found in parquet columns.") + action_matrix = np.stack(data[action_col]) + action_padding_ref = slice_list[0][3] + prepared_slices = [] + for a_slice, state_col, s_slice, action_padding, state_padding in slice_list: + if state_col not in data.columns: + raise ValueError(f"{state_col} not found in parquet columns.") + state_matrix = np.stack(data[state_col]) + state_part_full = state_matrix[:, s_slice[0] : s_slice[1]] + prepared_slices.append((a_slice, state_part_full, state_padding)) + for base_index in range(trajectory_length): + action_steps = np.array(action_indices) + base_index + action_chunk_full = _get_chunk(action_matrix, action_steps, action_padding_ref) + + for a_slice, state_part_full, state_padding in prepared_slices: + action_part_chunk = action_chunk_full[:, a_slice[0] : a_slice[1]] + state_chunk = _get_chunk(state_part_full, np.array(state_indices) + base_index, state_padding) + if action_part_chunk.shape[1] != state_chunk.shape[1]: + raise ValueError(f"Action/state dim mismatch for {action_col}:{a_slice}") + + out = relative_values( + action_part_chunk, + state_chunk[0], + is_quaternion=_action_slice_is_quaternion( + lerobot_modality_meta, a_slice + ), + ) + action_chunk_full[:, a_slice[0] : a_slice[1]] = out + + accum[action_col].append(action_chunk_full) + + rel_stats = copy.deepcopy(base_stats) + for action_col, series_list in accum.items(): + if not series_list: + continue + all_values = np.concatenate(series_list, axis=0).astype(np.float32) + rel_stats[action_col] = { + "mean": np.mean(all_values, axis=0).tolist(), + "std": np.std(all_values, axis=0).tolist(), + "min": np.min(all_values, axis=0).tolist(), + "max": np.max(all_values, axis=0).tolist(), + "q01": np.quantile(all_values, 0.01, axis=0).tolist(), + "q99": np.quantile(all_values, 0.99, axis=0).tolist(), + } + return rel_stats + +class ModalityConfig(BaseModel): + """Configuration for a modality.""" + + delta_indices: list[int] + """Delta indices to sample relative to the current index. The returned data will correspond to the original data at a sampled base index + delta indices.""" + modality_keys: list[str] + """The keys to load for the modality in the dataset.""" + + +class LeRobotSingleDataset(Dataset): + """ + Base dataset class for LeRobot that supports sharding. + """ + def __init__( + self, + dataset_path: Path | str, + modality_configs: dict[str, ModalityConfig], + embodiment_tag: str | EmbodimentTag, + video_backend: str = "decord", + video_backend_kwargs: dict | None = None, + transforms: ComposedModalityTransform | None = None, + delete_pause_frame: bool = False, + data_cfg = None, + **kwargs, + ): + """ + Initialize the dataset. + + Args: + dataset_path (Path | str): The path to the dataset. + modality_configs (dict[str, ModalityConfig]): The configuration for each modality. The keys are the modality names, and the values are the modality configurations. + See `ModalityConfig` for more details. + video_backend (str): Backend for video reading. + video_backend_kwargs (dict): Keyword arguments for the video backend when initializing the video reader. + transforms (ComposedModalityTransform): The transforms to apply to the dataset. + embodiment_tag (EmbodimentTag): Overload the embodiment tag for the dataset. e.g. define it as "new_embodiment" + """ + # first check if the path directory exists + self.data_cfg = data_cfg + if not Path(dataset_path).exists(): + raise FileNotFoundError(f"Dataset path {dataset_path} does not exist") + # indict letobot version + self._lerobot_version = self.data_cfg.get("lerobot_version", "v2.0") #self._indict_lerobot_version(**kwargs) + + self._action_mode = None + self._action_mode_state_map = {} + self._action_mode_apply_keys = None + + self.delete_pause_frame = delete_pause_frame + + self.modality_configs = modality_configs + self.video_backend = video_backend + self.video_backend_kwargs = video_backend_kwargs if video_backend_kwargs is not None else {} + self.transforms = ( + transforms if transforms is not None else ComposedModalityTransform(transforms=[]) + ) + + self._dataset_path = Path(dataset_path) + self._dataset_name = kwargs.get("dataset_name", self._dataset_path.name) + if isinstance(embodiment_tag, EmbodimentTag): + self.tag = embodiment_tag.value + else: + self.tag = embodiment_tag + + self._init_action_mode() + self._metadata = self._get_metadata(EmbodimentTag(self.tag)) + + # LeRobot-specific config + self._lerobot_modality_meta = self._get_lerobot_modality_meta() + self._lerobot_info_meta = self._get_lerobot_info_meta() + self._data_path_pattern = self._get_data_path_pattern() + self._video_path_pattern = self._get_video_path_pattern() + self._chunk_size = self._get_chunk_size() + self._tasks = self._get_tasks() + # self._episodes = self._get_episode_info() # TODO why we need this func + self.curr_traj_data = None + self.curr_traj_id = None + + self._trajectory_ids, self._trajectory_lengths = self._get_trajectories() + self._modality_keys = self._get_modality_keys() + self._delta_indices = self._get_delta_indices() + self._all_steps = self._get_all_steps() + self.set_transforms_metadata(self.metadata) + self.set_epoch(0) + + if int(os.environ.get("RANK", "0")) == 0: + print(f"Initialized dataset {self.dataset_name} with {embodiment_tag}") + + # Check if the dataset is valid + self._check_integrity() + + @property + def dataset_path(self) -> Path: + """The path to the dataset that contains the METADATA_FILENAME file.""" + return self._dataset_path + + @property + def metadata(self) -> DatasetMetadata: + """The metadata for the dataset, loaded from metadata.json in the dataset directory""" + return self._metadata + + @property + def trajectory_ids(self) -> np.ndarray: + """The trajectory IDs in the dataset, stored as a 1D numpy array of strings.""" + return self._trajectory_ids + + @property + def trajectory_lengths(self) -> np.ndarray: + """The trajectory lengths in the dataset, stored as a 1D numpy array of integers. + The order of the lengths is the same as the order of the trajectory IDs. + """ + return self._trajectory_lengths + + @property + def all_steps(self) -> list[tuple[int, int]]: + """The trajectory IDs and base indices for all steps in the dataset. + Example: + self.trajectory_ids: [0, 1, 2] + self.trajectory_lengths: [3, 2, 4] + return: [ + ("traj_0", 0), ("traj_0", 1), ("traj_0", 2), + ("traj_1", 0), ("traj_1", 1), + ("traj_2", 0), ("traj_2", 1), ("traj_2", 2), ("traj_2", 3) + ] + """ + return self._all_steps + + @property + def modality_keys(self) -> dict: + """The modality keys for the dataset. The keys are the modality names, and the values are the keys for each modality. + + Example: { + "video": ["video.image_side_0", "video.image_side_1"], + "state": ["state.eef_position", "state.eef_rotation"], + "action": ["action.eef_position", "action.eef_rotation"], + "language": ["language.human.task"], + "timestamp": ["timestamp"], + "reward": ["reward"], + } + """ + return self._modality_keys + + @property + def delta_indices(self) -> dict[str, np.ndarray]: + """The delta indices for the dataset. The keys are the modality.key, and the values are the delta indices for each modality.key.""" + return self._delta_indices + + @property + def dataset_name(self) -> str: + """The name of the dataset.""" + return self._dataset_name + + @property + def lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """The metadata for the LeRobot dataset.""" + return self._lerobot_modality_meta + + @property + def lerobot_info_meta(self) -> dict: + """The metadata for the LeRobot dataset.""" + return self._lerobot_info_meta + + @property + def data_path_pattern(self) -> str: + """The path pattern for the LeRobot dataset.""" + return self._data_path_pattern + + @property + def video_path_pattern(self) -> str: + """The path pattern for the LeRobot dataset.""" + return self._video_path_pattern + + @property + def chunk_size(self) -> int: + """The chunk size for the LeRobot dataset.""" + return self._chunk_size + + @property + def tasks(self) -> pd.DataFrame: + """The tasks for the dataset.""" + return self._tasks + + def _get_metadata(self, embodiment_tag: EmbodimentTag) -> DatasetMetadata: + """Get the metadata for the dataset. + + Returns: + dict: The metadata for the dataset. + """ + + # 1. Modality metadata + modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME + assert ( + modality_meta_path.exists() + ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" + # 1.1. State and action modalities + simplified_modality_meta: dict[str, dict] = {} + with open(modality_meta_path, "r") as f: + le_modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + for modality in ["state", "action"]: + simplified_modality_meta[modality] = {} + le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr( + le_modality_meta, modality + ) + for subkey in le_state_action_meta: + state_action_dtype = np.dtype(le_state_action_meta[subkey].dtype) + if np.issubdtype(state_action_dtype, np.floating): + continuous = True + else: + continuous = False + simplified_modality_meta[modality][subkey] = { + "absolute": le_state_action_meta[subkey].absolute, + "rotation_type": le_state_action_meta[subkey].rotation_type, + "shape": [ + le_state_action_meta[subkey].end - le_state_action_meta[subkey].start + ], + "continuous": continuous, + } + + # 1.2. Video modalities + le_info_path = self.dataset_path / LE_ROBOT_INFO_FILENAME + assert ( + le_info_path.exists() + ), f"Please provide a {LE_ROBOT_INFO_FILENAME} file in {self.dataset_path}" + with open(le_info_path, "r") as f: + le_info = json.load(f) + simplified_modality_meta["video"] = {} + for new_key in le_modality_meta.video: + original_key = le_modality_meta.video[new_key].original_key + if original_key is None: + original_key = new_key + le_video_meta = le_info["features"][original_key] + height = le_video_meta["shape"][le_video_meta["names"].index("height")] + width = le_video_meta["shape"][le_video_meta["names"].index("width")] + # NOTE(FH): different lerobot dataset versions have different keys for the number of channels and fps + try: + channels = le_video_meta["shape"][le_video_meta["names"].index("channel")] + fps = le_video_meta["video_info"]["video.fps"] + except (ValueError, KeyError): + try: + channels = le_video_meta["info"]["video.channels"] + fps = le_video_meta["info"]["video.fps"] + except (ValueError, KeyError): + # Fallback for image-only datasets (e.g. VLA-Arena) that lack video_info + channels = 3 + fps = le_info.get("fps", 30) + simplified_modality_meta["video"][new_key] = { + "resolution": [width, height], + "channels": channels, + "fps": fps, + } + + + # 2. Dataset statistics + def is_main(): + return (not dist.is_initialized()) or dist.get_rank() == 0 + + action_mode = _normalize_action_mode(self.data_cfg.get("action_mode", "abs") if self.data_cfg else "abs") + + stats_path = self.dataset_path / LE_ROBOT_STATS_FILENAME + action_cfg = self.modality_configs.get("action") + state_cfg = self.modality_configs.get("state") + action_keys_full = list(action_cfg.modality_keys) if action_cfg else [] + state_keys_full = list(state_cfg.modality_keys) if state_cfg else [] + action_indices = list(action_cfg.delta_indices) if action_cfg else None + state_indices = list(state_cfg.delta_indices) if state_cfg else None + + apply_keys = _normalize_action_mode_apply_keys( + self.data_cfg.get("action_mode_apply_keys", None) if self.data_cfg else None, + action_keys_full, + ) + normalized_state_map = _normalize_action_mode_state_map( + self.data_cfg.get("action_mode_state_map", {}) if self.data_cfg else {} + ) + stats_cache_config = _build_stats_cache_config( + action_mode=action_mode, + ) + precomputed_path = ( + self.data_cfg.get("precomputed_statistics_path") + if self.data_cfg + else None + ) + precomputed_filename = ( + self.data_cfg.get("precomputed_statistics_filename") + if self.data_cfg + else None + ) + if precomputed_path: + with open(precomputed_path, "r", encoding="utf-8") as statistics_file: + statistics_payload = json.load(statistics_file) + statistics_key = self.data_cfg.precomputed_statistics_key_map[ + self.dataset_name + ] + dataset_statistics = statistics_payload[statistics_key] + le_statistics = { + "action": dataset_statistics["action"], + "observation.state": dataset_statistics["state"], + } + elif precomputed_filename: + if action_mode != "abs": + raise ValueError( + "precomputed_statistics_filename currently requires action_mode=abs" + ) + precomputed_path = self.dataset_path / precomputed_filename + with open(precomputed_path, "r", encoding="utf-8") as statistics_file: + le_statistics = json.load(statistics_file) + required_columns = { + le_modality_meta.get_key_meta(key).original_key + for key in action_keys_full + state_keys_full + } + le_statistics = { + key: le_statistics[key] for key in required_columns + } + else: + parquet_files = list(self.dataset_path.glob(LE_ROBOT_DATA_FILENAME)) + parquet_files_filtered = [ + pf for pf in parquet_files if "episode_033675.parquet" not in pf.name + ] + + if is_main(): + le_statistics = _load_or_compute_statistics( + stats_path, + stats_cache_config=stats_cache_config, + parquet_paths=parquet_files_filtered, + dataset_name=self.dataset_name, + action_mode=action_mode, + lerobot_modality_meta=le_modality_meta, + action_keys_full=action_keys_full, + state_keys_full=state_keys_full, + action_indices=action_indices, + state_indices=state_indices, + action_mode_apply_keys=apply_keys, + action_mode_state_map=normalized_state_map, + ) + else: + le_statistics = None + + if dist.is_initialized(): + dist.barrier() + + if le_statistics is None: + le_statistics = _load_stats_cache( + stats_path, + stats_cache_config, + invalidate_legacy=False, + ) + if le_statistics is None: + raise RuntimeError( + f"Dataset statistics cache is missing or invalid after sync: {stats_path}" + ) + + for stat in le_statistics.values(): + DatasetStatisticalValues.model_validate(stat) + + + dataset_statistics = {} + for our_modality in ["state", "action"]: + dataset_statistics[our_modality] = {} + for subkey in simplified_modality_meta[our_modality]: + dataset_statistics[our_modality][subkey] = {} + state_action_meta = le_modality_meta.get_key_meta(f"{our_modality}.{subkey}") + assert isinstance(state_action_meta, LeRobotStateActionMetadata) + le_modality = state_action_meta.original_key + for stat_name in le_statistics[le_modality]: + indices = np.arange( + state_action_meta.start, + state_action_meta.end, + ) + stat = np.array(le_statistics[le_modality][stat_name]) + dataset_statistics[our_modality][subkey][stat_name] = stat[indices].tolist() + + # 3. Full dataset metadata + metadata = DatasetMetadata( + statistics=dataset_statistics, # type: ignore + modalities=simplified_modality_meta, # type: ignore + embodiment_tag=embodiment_tag, + ) + + return metadata + + def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]: + """Get the trajectories in the dataset.""" + # Get trajectory lengths, IDs, and whitelist from dataset metadata + # v2.0 + if self._lerobot_version == "v2.0": + file_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME + with open(file_path, "r") as f: + episode_metadata = [json.loads(line) for line in f] + trajectory_ids = [] + trajectory_lengths = [] + for episode in episode_metadata: + trajectory_ids.append(episode["episode_index"]) + trajectory_lengths.append(episode["length"]) + return np.array(trajectory_ids), np.array(trajectory_lengths) + # v3.0 + elif self._lerobot_version == "v3.0": + file_paths = sorted(list((self.dataset_path).glob(LE_ROBOT3_EPISODE_FILENAME))) + trajectory_ids = [] + trajectory_lengths = [] + # data_chunck_index = [] + # data_file_index = [] + # vido_from_index = [] + self.trajectory_ids_to_metadata = {} + for file_path in file_paths: + episodes_data = pd.read_parquet(file_path) + timestamp_cols = [ + c + for c in episodes_data.columns + if str(c).startswith("videos/") and str(c).endswith("/from_timestamp") + ] + for index, episode in episodes_data.iterrows(): + trajectory_ids.append(episode["episode_index"]) + trajectory_lengths.append(episode["length"]) + + from_timestamps = {} + for col in timestamp_cols: + value = episode[col] + if pd.isna(value): + continue + # videos/{video_key}/from_timestamp -> {video_key} + video_key = str(col)[len("videos/") : -len("/from_timestamp")] + from_timestamps[video_key] = float(value) + + # Collect video file indices for each video key + video_file_indices = {} + for col in timestamp_cols: + video_key = str(col)[len("videos/") : -len("/from_timestamp")] + chunk_col = f"videos/{video_key}/chunk_index" + file_col = f"videos/{video_key}/file_index" + if chunk_col in episode and file_col in episode: + video_file_indices[video_key] = { + "chunk_index": int(episode[chunk_col]), + "file_index": int(episode[file_col]), + } + print(video_file_indices) + episode_meta = { + "data/chunk_index": episode["data/chunk_index"], + "data/file_index": episode["data/file_index"], + "data/file_from_index": index, + "videos/from_timestamps": from_timestamps, + "videos/file_indices": video_file_indices, + } + # episode_meta = { + # "data/chunk_index": episode["data/chunk_index"], + # "data/file_index": episode["data/file_index"], + # "data/file_from_index": index, + # "videos/from_timestamps": from_timestamps, + # } + self.trajectory_ids_to_metadata[trajectory_ids[-1]] = episode_meta + + # Should be able to directly read the saved index info here + return np.array(trajectory_ids), np.array(trajectory_lengths) + + def _get_all_steps(self) -> list[tuple[int, int]]: + """Get the trajectory IDs and base indices for all steps in the dataset. + + Returns: + list[tuple[str, int]]: A list of (trajectory_id, base_index) tuples. + """ + def is_main(): + return (not dist.is_initialized()) or dist.get_rank() == 0 + + config_key = self._get_steps_config_key() + steps_filename = "steps_data_index.pkl" + steps_path = self.dataset_path / "meta" / steps_filename + + # ---------- try to read from cache ---------- + if steps_path.exists(): + try: + with open(steps_path, "rb") as f: + cached_data = pickle.load(f) + if cached_data.get("config_key") != config_key: + raise ValueError( + "cached step index was built for a different dataset configuration" + ) + if cached_data.get("num_trajectories") != len(self.trajectory_ids): + raise ValueError("cached step index has a different trajectory count") + steps = cached_data["steps"] + if cached_data.get("total_steps") != len(steps): + raise ValueError("cached step index has an inconsistent step count") + return steps + except Exception as e: + # include EOFError / PickleError / KeyError + print( + f"[RANK {os.environ.get('RANK', 'NA')}] " + f"Failed to load cached steps ({e}), will rebuild." + ) + + # ---------- only build by rank0 ---------- + if is_main(): + all_steps = self._get_all_steps_single_process() + + cache_data = { + "config_key": config_key, + "steps": all_steps, + "num_trajectories": len(self.trajectory_ids), + "total_steps": len(all_steps), + "computed_timestamp": pd.Timestamp.now().isoformat(), + "delete_pause_frame": self.delete_pause_frame, + } + + steps_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = steps_path.with_suffix(".tmp") + + with open(tmp_path, "wb") as f: + pickle.dump(cache_data, f, protocol=pickle.HIGHEST_PROTOCOL) + os.replace(tmp_path, steps_path) + + print(f"[RANK 0] Cached steps saved to {steps_path}") + + # ---------- sync after rank0 ---------- + if dist.is_initialized(): + dist.barrier() + + # ---------- read by all rank ---------- + with open(steps_path, "rb") as f: + cached_data = pickle.load(f) + + return cached_data["steps"] + + def _get_steps_config_key(self) -> str: + """Generate a configuration key for steps caching.""" + config_dict = { + "delete_pause_frame": self.delete_pause_frame, + "dataset_name": self.dataset_name, + "lerobot_version": self._lerobot_version, + "language_keys": self.modality_keys.get("language", []), + "trajectory_ids": self.trajectory_ids.tolist(), + "trajectory_lengths": self.trajectory_lengths.tolist(), + } + # Create a hash of the configuration + config_str = str(sorted(config_dict.items())) + return hashlib.md5(config_str.encode()).hexdigest()[:12] # + + + def _get_all_steps_single_process(self) -> list[tuple[int, int]]: + """Original single-process implementation as fallback.""" + all_steps: list[tuple[int, int]] = [] + skipped_trajectories = 0 + processed_trajectories = 0 + + # Check if language modality is configured + has_language_modality = 'language' in self.modality_keys and len(self.modality_keys['language']) > 0 + # TODO why trajectory_length here, why not use data length? + for trajectory_id, trajectory_length in tqdm(zip(self.trajectory_ids, self.trajectory_lengths), total=len(self.trajectory_ids), desc="Getting All Step"): + try: + if self._lerobot_version == "v2.0": + data = self.get_trajectory_data(trajectory_id) + elif self._lerobot_version == "v3.0": + data = self.get_trajectory_data_lerobot_v3(trajectory_id) + + trajectory_skipped = False + + # Check if trajectory has valid language instruction (if language modality is configured) + if has_language_modality: + self.curr_traj_data = data # Set current trajectory data for get_language to work + + language_instruction = self.get_language(trajectory_id, self.modality_keys['language'][0], 0) + if not language_instruction or language_instruction[0] == "": + print(f"Skipping trajectory {trajectory_id} due to empty language instruction") + skipped_trajectories += 1 + trajectory_skipped = True + continue + + except Exception as e: + print(f"Skipping trajectory {trajectory_id} due to read error: {e}") + skipped_trajectories += 1 + trajectory_skipped = True + continue + + if not trajectory_skipped: + processed_trajectories += 1 + + for base_index in range(trajectory_length): + all_steps.append((trajectory_id, base_index)) + + # Print summary statistics + print(f"Single-process summary: Processed {processed_trajectories} trajectories, skipped {skipped_trajectories} empty trajectories") + print(f"Total steps: {len(all_steps)} from {len(self.trajectory_ids)} trajectories") + + return all_steps + + def _get_position_and_gripper_values(self, data: pd.DataFrame) -> tuple[list, list]: + """Get position and gripper values based on available columns in the dataset.""" + # Get action keys from modality_keys + action_keys = self.modality_keys.get('action', []) + + # Extract position data + delta_position_values = None + position_candidates = ['delta_eef_position'] + coordinate_candidates = ['x', 'y', 'z'] + + # First try combined position fields + for pos_key in position_candidates: + full_key = f"action.{pos_key}" + if full_key in action_keys: + try: + # Get the lerobot key for this modality + le_action_cfg = self.lerobot_modality_meta.action + subkey = pos_key + if subkey in le_action_cfg: + le_key = le_action_cfg[subkey].original_key or subkey + if le_key in data.columns: + data_array = np.stack(data[le_key]) + le_indices = np.arange(le_action_cfg[subkey].start, le_action_cfg[subkey].end) + filtered_data = data_array[:, le_indices] + delta_position_values = filtered_data.tolist() + break + except Exception: + continue + + # If combined fields not found, try individual x,y,z coordinates + if delta_position_values is None: + x_data, y_data, z_data = None, None, None + for coord in coordinate_candidates: + full_key = f"action.{coord}" + if full_key in action_keys: + try: + le_action_cfg = self.lerobot_modality_meta.action + if coord in le_action_cfg: + le_key = le_action_cfg[coord].original_key or coord + if le_key in data.columns: + data_array = np.stack(data[le_key]) + le_indices = np.arange(le_action_cfg[coord].start, le_action_cfg[coord].end) + coord_data = data_array[:, le_indices].flatten() + if coord == 'x': + x_data = coord_data + elif coord == 'y': + y_data = coord_data + elif coord == 'z': + z_data = coord_data + except Exception: + continue + + if x_data is not None and y_data is not None and z_data is not None: + delta_position_values = np.column_stack((x_data, y_data, z_data)).tolist() + + if delta_position_values is None: + # Fallback to the old hardcoded approach if metadata approach fails + if 'action.delta_eef_position' in data.columns: + delta_position_values = data['action.delta_eef_position'].to_numpy().tolist() + elif all(col in data.columns for col in ['action.x', 'action.y', 'action.z']): + x_vals = data['action.x'].to_numpy() + y_vals = data['action.y'].to_numpy() + z_vals = data['action.z'].to_numpy() + delta_position_values = np.column_stack((x_vals, y_vals, z_vals)).tolist() + else: + raise ValueError(f"No suitable position columns found. Available columns: {data.columns.tolist()}") + + # Extract gripper data + gripper_values = None + gripper_candidates = ['gripper_close', 'gripper'] + + for grip_key in gripper_candidates: + full_key = f"action.{grip_key}" + if full_key in action_keys: + try: + le_action_cfg = self.lerobot_modality_meta.action + if grip_key in le_action_cfg: + le_key = le_action_cfg[grip_key].original_key or grip_key + if le_key in data.columns: + data_array = np.stack(data[le_key]) + le_indices = np.arange(le_action_cfg[grip_key].start, le_action_cfg[grip_key].end) + gripper_data = data_array[:, le_indices].flatten() + gripper_values = gripper_data.tolist() + break + except Exception: + continue + + if gripper_values is None: + # Fallback to the old hardcoded approach if metadata approach fails + if 'action.gripper_close' in data.columns: + gripper_values = data['action.gripper_close'].to_numpy().tolist() + elif 'action.gripper' in data.columns: + gripper_values = data['action.gripper'].to_numpy().tolist() + else: + raise ValueError(f"No suitable gripper columns found. Available columns: {data.columns.tolist()}") + + return delta_position_values, gripper_values + + def _get_modality_keys(self) -> dict: + """Get the modality keys for the dataset. + The keys are the modality names, and the values are the keys for each modality. + See property `modality_keys` for the expected format. + """ + modality_keys = defaultdict(list) + for modality, config in self.modality_configs.items(): + modality_keys[modality] = config.modality_keys + return modality_keys + + def _get_delta_indices(self) -> dict[str, np.ndarray]: + """Restructure the delta indices to use modality.key as keys instead of just the modalities.""" + delta_indices: dict[str, np.ndarray] = {} + for config in self.modality_configs.values(): + for key in config.modality_keys: + delta_indices[key] = np.array(config.delta_indices) + return delta_indices + + def _init_action_mode(self) -> None: + if self.data_cfg is None: + self._action_mode = "abs" + return + + action_mode = self.data_cfg.get("action_mode", "abs") + if action_mode is None: + action_mode = "abs" + action_mode = _normalize_action_mode(action_mode) + if action_mode not in {"abs", "delta", "rel"}: + raise ValueError(f"Invalid action_mode: {action_mode}. Expected one of: abs, delta, rel.") + self._action_mode = action_mode + + apply_keys = _normalize_action_mode_apply_keys(self.data_cfg.get("action_mode_apply_keys", None)) + if apply_keys: + self._action_mode_apply_keys = apply_keys + + self._action_mode_state_map = _normalize_action_mode_state_map( + self.data_cfg.get("action_mode_state_map", {}) or {} + ) + + def _action_key_is_quaternion(self, action_key: str) -> bool: + if not action_key.startswith("action."): + return False + subkey = action_key.replace("action.", "", 1) + try: + cfg = self._lerobot_modality_meta.action[subkey] + except (AttributeError, KeyError): + return False + return rotation_type_is_quaternion(cfg.rotation_type) + + def _infer_state_key_for_action(self, action_key: str) -> str | None: + if action_key in self._action_mode_state_map: + return self._action_mode_state_map[action_key] + + if not action_key.startswith("action."): + return None + base = action_key.replace("action.", "", 1) + if f"state.{base}" in self.modality_keys.get("state", []): + return f"state.{base}" + return None + + def _apply_action_mode(self, data: dict) -> dict: + if self._action_mode in (None, "abs"): + return data + + action_keys = self._action_mode_apply_keys or self.modality_keys.get("action", []) + for action_key in action_keys: + if action_key not in data: + print(f"[WARNING] Action key {action_key} not found in data") + continue + state_key = self._infer_state_key_for_action(action_key) + + # for safety, check if the state key is valid + if state_key is None or state_key not in data: + continue + + action_values = np.asarray(data[action_key]) + state_values = np.asarray(data[state_key]) + if action_values.ndim != 2 or state_values.ndim != 2: + raise ValueError( + f"Expected 2D arrays for action/state, got {action_key}: {action_values.shape}, {state_key}: {state_values.shape}" + ) + if action_values.shape[1] != state_values.shape[1]: + raise ValueError( + f"Action/state dim mismatch for {action_key} vs {state_key}: {action_values.shape} vs {state_values.shape}" + ) + + state0 = state_values[0] + is_quaternion = self._action_key_is_quaternion(action_key) + if self._action_mode == "delta": + out = action_values.copy() + if is_quaternion: + raise ValueError( + f"action_mode=delta is not defined for quaternion key {action_key}; " + "use rel (SE(3) relative to the current pose) or abs." + ) + if len(out) > 1: + out[1:] = action_values[1:] - action_values[:-1] + out[0] = action_values[0] - state0 + elif self._action_mode == "rel": + out = relative_values(action_values, state0, is_quaternion=is_quaternion) + else: + out = action_values + + data[action_key] = out + + return data + + def _get_lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """Get the metadata for the LeRobot dataset.""" + modality_meta_path = self.dataset_path / LE_ROBOT_MODALITY_FILENAME + assert ( + modality_meta_path.exists() + ), f"Please provide a {LE_ROBOT_MODALITY_FILENAME} file in {self.dataset_path}" + with open(modality_meta_path, "r") as f: + modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + return modality_meta + + def _get_lerobot_info_meta(self) -> dict: + """Get the metadata for the LeRobot dataset.""" + info_meta_path = self.dataset_path / LE_ROBOT_INFO_FILENAME + with open(info_meta_path, "r") as f: + info_meta = json.load(f) + return info_meta + + def _get_data_path_pattern(self) -> str: + """Get the data path pattern for the LeRobot dataset.""" + return self.lerobot_info_meta["data_path"] + + def _get_video_path_pattern(self) -> str: + """Get the video path pattern for the LeRobot dataset.""" + return self.lerobot_info_meta["video_path"] + + def _get_chunk_size(self) -> int: + """Get the chunk size for the LeRobot dataset.""" + return self.lerobot_info_meta["chunks_size"] + + def _get_tasks(self) -> pd.DataFrame: + """Get the tasks for the dataset.""" + if self._lerobot_version == "v2.0": + tasks_path = self.dataset_path / LE_ROBOT_TASKS_FILENAME + with open(tasks_path, "r") as f: + tasks = [json.loads(line) for line in f] + df = pd.DataFrame(tasks) + return df.set_index("task_index") + + elif self._lerobot_version == "v3.0": + tasks_path = self.dataset_path / LE_ROBOT3_TASKS_FILENAME + df = pd.read_parquet(tasks_path) + df = df.reset_index() # convert index to a column, typically named 'index' + df = df.rename(columns={'index': 'task'}) # rename 'index' column to 'task' + df = df[['task_index', 'task']] # reorder columns + return df + def _check_integrity(self): + """Use the config to check if the keys are valid and detect silent data corruption.""" + ERROR_MSG_HEADER = f"Error occurred in initializing dataset {self.dataset_name}:\n" + + for modality_config in self.modality_configs.values(): + for key in modality_config.modality_keys: + if key == "lapa_action" or key == "dream_actions": + continue # no need for any metadata for lapa actions because it comes normalized + # Check if the key is valid + try: + self.lerobot_modality_meta.get_key_meta(key) + except Exception as e: + raise ValueError( + ERROR_MSG_HEADER + f"Unable to find key {key} in modality metadata:\n{e}" + ) + + def set_transforms_metadata(self, metadata: DatasetMetadata): + """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" + self.transforms.set_metadata(metadata) + + def set_epoch(self, epoch: int): + """Set the epoch for the dataset. + + Args: + epoch (int): The epoch to set. + """ + self.epoch = epoch + + def __len__(self) -> int: + """Get the total number of data points in the dataset. + + Returns: + int: the total number of data points in the dataset. + """ + return len(self.all_steps) + + def __str__(self) -> str: + """Get the description of the dataset.""" + return f"{self.dataset_name} ({len(self)} steps)" + + + def __getitem__(self, index: int) -> dict: + """Get the data for a single step in a trajectory. + + Args: + index (int): The index of the step to get. + + Returns: + dict: The data for the step. + """ + trajectory_id, base_index = self.all_steps[index] + raw_data = self.get_step_data(trajectory_id, base_index) + data = self.transforms(raw_data) + return self._pack_sample(data) + + def _pack_sample(self, data: dict) -> dict: + """Pack transformed modality data into training sample format.""" + step_images = [] + for video_key in self.modality_keys["video"]: + image = data[video_key][0] + image = Image.fromarray(image).resize((224, 224)) + step_images.append(image) + + language = data[self.modality_keys["language"][0]][0] + action = [] + for action_key in self.modality_keys["action"]: + action.append(data[action_key]) + action = np.concatenate(action, axis=1).astype(np.float16) + + sample = { + "action": action, + "image": step_images, + "lang": language, + "robot_tag": self.tag + } + + if self.data_cfg is not None and self.data_cfg.get("include_state", False) not in ["False", False]: + state = [] + for state_key in self.modality_keys.get("state", []): + state.append(data[state_key]) + if not state: + import warnings + warnings.warn( + "include_state=True but no state modality keys found in modality_configs " + "(state modality may be disabled in the dataset's DataConfig). " + "Skipping state packing.", + stacklevel=2, + ) + else: + state = np.concatenate(state, axis=1).astype(np.float16) + sample["state"] = state + + return sample + + def get_step_data(self, trajectory_id: int, base_index: int) -> dict: + """Get the RAW data for a single step in a trajectory. No transforms are applied. + + Args: + trajectory_id (int): The name of the trajectory. + base_index (int): The base step index in the trajectory. + + Returns: + dict: The RAW data for the step. + + Example return: + { + "video": { + "video.image_side_0": [B, T, H, W, C], + "video.image_side_1": [B, T, H, W, C], + }, + "state": { + "state.eef_position": [B, T, state_dim], + "state.eef_rotation": [B, T, state_dim], + }, + "action": { + "action.eef_position": [B, T, action_dim], + "action.eef_rotation": [B, T, action_dim], + }, + } + """ + data = {} + # Get the data for all modalities # just for action base data + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + # Read the requested trajectory from the active trajectory table. + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index) + data = self._apply_action_mode(data) + return data + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the data for a trajectory.""" + if self._lerobot_version == "v2.0": + + if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: + return self.curr_traj_data + else: + chunk_index = self.get_episode_chunk(trajectory_id) + parquet_path = self.dataset_path / self.data_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id + ) + assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" + return pd.read_parquet(parquet_path) + elif self._lerobot_version == "v3.0": + return self.get_trajectory_data_lerobot_v3(trajectory_id) + + def get_trajectory_data_lerobot_v3(self, trajectory_id: int) -> pd.DataFrame: + """Get the data for a trajectory from lerobot v3.""" + if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: + return self.curr_traj_data + else: #TODO check detail later + episode_meta = self.trajectory_ids_to_metadata[trajectory_id] + chunk_index = episode_meta["data/chunk_index"] + file_index = self.get_episode_file_index(trajectory_id) + # file_from_index = self.get_episode_file_from_index(trajectory_id) + + + parquet_path = self.dataset_path / self.data_path_pattern.format( + chunk_index=chunk_index, file_index=file_index + ) + assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" + file_data = pd.read_parquet(parquet_path) + + # filter by trajectory_id + episode_data = file_data.loc[file_data["episode_index"] == trajectory_id].copy() + return episode_data + + + def get_trajectory_index(self, trajectory_id: int) -> int: + """Get the index of the trajectory in the dataset by the trajectory ID. + This is useful when you need to get the trajectory length or sampling weight corresponding to the trajectory ID. + + Args: + trajectory_id (str): The ID of the trajectory. + + Returns: + int: The index of the trajectory in the dataset. + """ + trajectory_indices = np.where(self.trajectory_ids == trajectory_id)[0] + if len(trajectory_indices) != 1: + raise ValueError( + f"Error finding trajectory index for {trajectory_id}, found {trajectory_indices=}" + ) + return trajectory_indices[0] + + def get_episode_chunk(self, ep_index: int) -> int: + """Get the chunk index for an episode index.""" + return ep_index // self.chunk_size + def get_episode_file_index(self, ep_index: int) -> int: + """Get the file index for an episode index.""" + episode_meta = self.trajectory_ids_to_metadata[ep_index] + return episode_meta["data/file_index"] + + def get_episode_file_from_index(self, ep_index: int) -> int: + """Get the file from index for an episode index.""" + episode_meta = self.trajectory_ids_to_metadata[ep_index] + return episode_meta["data/file_from_index"] + + + def retrieve_data_and_pad( + self, + array: np.ndarray, + step_indices: np.ndarray, + max_length: int, + padding_strategy: str = "first_last", + ) -> np.ndarray: + """Retrieve the data from the dataset and pad it if necessary. + Args: + array (np.ndarray): The array to retrieve the data from. + step_indices (np.ndarray): The step indices to retrieve the data for. + max_length (int): The maximum length of the data. + padding_strategy (str): The padding strategy, either "first" or "last". + """ + # Get the padding indices + front_padding_indices = step_indices < 0 + end_padding_indices = step_indices >= max_length + padding_positions = np.logical_or(front_padding_indices, end_padding_indices) + # Retrieve the data with the non-padding indices + # If there exists some padding, Given T step_indices, the shape of the retrieved data will be (T', ...) where T' < T + raw_data = array[step_indices[~padding_positions]] + assert isinstance(raw_data, np.ndarray), f"{type(raw_data)=}" + # This is the shape of the output, (T, ...) + if raw_data.ndim == 1: + expected_shape = (len(step_indices),) + else: + expected_shape = (len(step_indices), *array.shape[1:]) + + # Pad the data + output = np.zeros(expected_shape) + # Assign the non-padded data + output[~padding_positions] = raw_data + # If there exists some padding, pad the data + if padding_positions.any(): + if padding_strategy == "first_last": + # Use first / last step data to pad + front_padding_data = array[0] + end_padding_data = array[-1] + output[front_padding_indices] = front_padding_data + output[end_padding_indices] = end_padding_data + elif padding_strategy == "zero": + # Use zero padding + output[padding_positions] = 0 + else: + raise ValueError(f"Invalid padding strategy: {padding_strategy}") + return output + + def get_video_path(self, trajectory_id: int, key: str) -> Path: + chunk_index = self.get_episode_chunk(trajectory_id) + original_key = self.lerobot_modality_meta.video[key].original_key + if original_key is None: + original_key = key + if self._lerobot_version == "v2.0": + video_filename = self.video_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key + ) + elif self._lerobot_version == "v3.0": + episode_meta = self.trajectory_ids_to_metadata[trajectory_id] + + video_file_indices = episode_meta.get("videos/file_indices", {}) + if original_key in video_file_indices: + video_chunk_index = video_file_indices[original_key]["chunk_index"] + video_file_index = video_file_indices[original_key]["file_index"] + else: + video_chunk_index = episode_meta["data/chunk_index"] + video_file_index = episode_meta["data/file_index"] + video_filename = self.video_path_pattern.format( + video_key=original_key, + chunk_index=episode_meta["data/chunk_index"], + file_index=episode_meta["data/file_index"], + ) + return self.dataset_path / video_filename + + def get_video( + self, + trajectory_id: int, + key: str, + base_index: int, + ) -> np.ndarray: + """Get the video frames for a trajectory by a base index. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (str): The ID of the trajectory. + key (str): The key of the video. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) + """ + # Get the step indices + step_indices = self.delta_indices[key] + base_index + # print(f"{step_indices=}") + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Get the sub-key + key = key.replace("video.", "") + + # Image-only LeRobot datasets (e.g. VLA-Arena) may store frames directly + # in parquet columns and have total_videos == 0 (no mp4 files). In this + # case, load frames from the original image column and pad by step indices. + original_key = self.lerobot_modality_meta.video[key].original_key + if original_key is None: + original_key = key + if self.curr_traj_data is not None and original_key in self.curr_traj_data.columns: + image_entries = self.curr_traj_data[original_key].tolist() + + def _decode_image_entry(entry): + if isinstance(entry, np.ndarray): + return entry + if isinstance(entry, Image.Image): + return np.array(entry) + if isinstance(entry, dict): + img_bytes = entry.get("bytes", None) + img_path = entry.get("path", None) + + if img_bytes is not None: + return np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB")) + + if img_path is not None: + path_obj = Path(img_path) + if not path_obj.is_absolute(): + path_obj = self.dataset_path / path_obj + return np.array(Image.open(path_obj).convert("RGB")) + + raise TypeError(f"Unsupported image entry type: {type(entry)}") + + frames = [] + for idx in step_indices: + safe_idx = int(min(max(idx, 0), len(image_entries) - 1)) + frames.append(_decode_image_entry(image_entries[safe_idx])) + + return np.stack(frames) + + video_path = self.get_video_path(trajectory_id, key) + # Get the action/state timestamps for each frame in the video + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert "timestamp" in self.curr_traj_data.columns, f"No timestamp found in {trajectory_id=}" + timestamp: np.ndarray = self.curr_traj_data["timestamp"].to_numpy() + # Get the corresponding video timestamps from the step indices + video_timestamp = timestamp[step_indices] + if self._lerobot_version == "v3.0": + episode_meta = self.trajectory_ids_to_metadata.get(trajectory_id, {}) + from_timestamps = episode_meta.get("videos/from_timestamps", {}) + original_video_key = self.lerobot_modality_meta.video[key].original_key + if original_video_key is None: + original_video_key = key + from_timestamp = float(from_timestamps.get(original_video_key, 0.0)) + video_timestamp = video_timestamp + from_timestamp + + return get_frames_by_timestamps( + video_path.as_posix(), + video_timestamp, + video_backend=self.video_backend, # TODO + video_backend_kwargs=self.video_backend_kwargs, + ) + + def get_state_or_action( + self, + trajectory_id: int, + modality: str, + key: str, + base_index: int, + ) -> np.ndarray: + """Get the state or action data for a trajectory by a base index. + If the step indices are out of range, pad with the data: + if the data is stored in absolute format, pad with the first or last step data; + otherwise, pad with zero. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + + Returns: + np.ndarray: The data for the trajectory and step indices. + """ + # Get the step indices + step_indices = self.delta_indices[key] + base_index + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + key = key.replace(modality + ".", "") + # Get the lerobot key + le_state_or_action_cfg = getattr(self.lerobot_modality_meta, modality) + le_key = le_state_or_action_cfg[key].original_key + if le_key is None: + le_key = key + # Get the data array, shape: (T, D) + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + assert le_key in self.curr_traj_data.columns, f"No {le_key} found in {trajectory_id=}" + data_array: np.ndarray = np.stack(self.curr_traj_data[le_key]) # type: ignore + assert data_array.ndim == 2, f"Expected 2D array, got key {le_key} is{data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[key].start, + le_state_or_action_cfg[key].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[key] + + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + # padding_strategy="zero", # HACK for realdata + ) + + def get_language( + self, + trajectory_id: int, + key: str, + base_index: int, + ) -> list[str]: + """Get the language annotation data for a trajectory by step indices. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + key (str): The key of the annotation. + base_index (int): The base index of the trajectory. + + Returns: + list[str]: The annotation data for the trajectory and step indices. If no matching data is found, return empty strings. + """ + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + # Get the step indices + step_indices = self.delta_indices[key] + base_index + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + # Get the end times corresponding to the closest indices + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, max_length - 1) + # Get the annotations + task_indices: list[int] = [] + assert key.startswith( + "annotation." + ), f"Language key must start with 'annotation.', got {key}" + subkey = key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + assert annotation_meta is not None, f"Annotation metadata is None for {subkey}" + assert ( + subkey in annotation_meta + ), f"Annotation key {subkey} not found in metadata, available annotation keys: {annotation_meta.keys()}" + subkey_meta = annotation_meta[subkey] + original_key = subkey_meta.original_key + if original_key is None: + original_key = key + for i in range(len(step_indices)): # + # task_indices.append(self.curr_traj_data[original_key][step_indices[i]].item()) + value = self.curr_traj_data[original_key].iloc[step_indices[i]] # TODO check v2.0 + task_indices.append(value if isinstance(value, (int, float)) else value.item()) + + return self.tasks.loc[task_indices]["task"].tolist() + + def get_data_by_modality( + self, + trajectory_id: int, + modality: str, + key: str, + base_index: int, + ): + """Get the data corresponding to the modality for a trajectory by a base index. + This method will call the corresponding helper method based on the modality. + See the helper methods for more details. + NOTE: For the language modality, the data is padded with empty strings if no matching data is found. + + Args: + dataset (BaseSingleDataset): The dataset to retrieve the data from. + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data. + key (str): The key of the data. + base_index (int): The base index of the trajectory. + """ + if modality == "video": + return self.get_video(trajectory_id, key, base_index) + elif modality == "state" or modality == "action": + return self.get_state_or_action(trajectory_id, modality, key, base_index) + elif modality == "language": + return self.get_language(trajectory_id, key, base_index) + else: + raise ValueError(f"Invalid modality: {modality}") + + def _save_dataset_statistics_(self, save_path: Path | str, format: str = "json") -> None: + """ + Save dataset statistics to specified path in the required format. + Only includes statistics for keys that are actually used in the dataset. + Gripper-related keys will be placed at the end. + + Args: + save_path (Path | str): Path to save the statistics file + format (str): Save format, currently only supports "json" + """ + save_path = Path(save_path) + save_path.parent.mkdir(parents=True, exist_ok=True) + + # Build the data structure to save + statistics_data = {} + + # Get used modality keys + used_action_keys, used_state_keys = get_used_modality_keys(self.modality_keys) + + # Organize statistics by tag + tag = self.tag + tag_stats = {} + + # Process action statistics (only for used keys) + if hasattr(self.metadata.statistics, 'action') and self.metadata.statistics.action: + action_stats = self.metadata.statistics.action + + # Filter to only include used action keys and reorder: non-gripper first, gripper last + non_gripper_keys = [] + gripper_keys = [] + + for key in action_stats.keys(): + if key in used_action_keys: + if "gripper" in key.lower(): + gripper_keys.append(key) + else: + non_gripper_keys.append(key) + + # Reorder: non-gripper first, gripper last + reordered_keys = non_gripper_keys + gripper_keys + + filtered_action_stats = {} + for key in reordered_keys: + filtered_action_stats[key] = action_stats[key] + + if filtered_action_stats: + # Combine statistics from filtered action sub-keys + combined_action_stats = combine_modality_stats(filtered_action_stats) + + # mask=False for dimensions whose normalization mode is "binary" + _action_norm_modes = _extract_action_normalization_modes(self.transforms) + mask = generate_action_mask_for_used_keys( + self.metadata.modalities.action, filtered_action_stats.keys(), + normalization_modes=_action_norm_modes, + ) + combined_action_stats["mask"] = mask + + tag_stats["action"] = combined_action_stats + + # Process state statistics (only for used keys) + if hasattr(self.metadata.statistics, 'state') and self.metadata.statistics.state: + state_stats = self.metadata.statistics.state + + # Filter to only include used state keys, optionally reorder gripper to end + non_gripper_keys = [] + gripper_keys = [] + + for key in state_stats.keys(): + if key in used_state_keys: + if "gripper" in key.lower(): + gripper_keys.append(key) + else: + non_gripper_keys.append(key) + + # Reorder: non-gripper first, gripper last + reordered_keys = non_gripper_keys + gripper_keys + + filtered_state_stats = {} + for key in reordered_keys: + filtered_state_stats[key] = state_stats[key] + + if filtered_state_stats: + combined_state_stats = combine_modality_stats(filtered_state_stats) + tag_stats["state"] = combined_state_stats + + # Add dataset counts + tag_stats["num_transitions"] = len(self) + tag_stats["num_trajectories"] = len(self.trajectory_ids) + + statistics_data[tag] = tag_stats + + # Save as JSON file + if format.lower() == "json": + if not str(save_path).endswith('.json'): + save_path = save_path.with_suffix('.json') + with open(save_path, 'w', encoding='utf-8') as f: + json.dump(statistics_data, f, indent=2, ensure_ascii=False) + else: + raise ValueError(f"Unsupported format: {format}. Currently only 'json' is supported.") + + print(f"Single dataset statistics saved to: {save_path}") + print(f"Used action keys (reordered): {list(used_action_keys)}") + print(f"Used state keys (reordered): {list(used_state_keys)}") + + +class CachedLeRobotSingleDataset(LeRobotSingleDataset): + def __init__(self, img_resize: tuple[int, int] | None = None, *args, **kwargs): + """ + This class caches the video frames for each trajectory and key. + It is recommended to use this class if the video frames need to be accessed multiple times. + + Args: + resize_img (tuple[int, int], optional): The size to resize the video frames to reduce memory usage. + """ + # Convert img_resize to tuple if it is not already + if img_resize is not None and not isinstance(img_resize, tuple): + img_resize = tuple(img_resize) + assert len(img_resize) == 2, f"Expected tuple of length 2, got {img_resize}" + self.img_resize = img_resize + + # Initialize img_resize attribute first to ensure it exists + super().__init__(*args, **kwargs) + cached_frames: dict[str, np.ndarray] = {} + + for key in self.modality_keys["video"]: + all_frames = [] + original_key = key + key = key.replace("video.", "") + for trajectory_id, trajectory_length in tqdm( + zip(self.trajectory_ids, self.trajectory_lengths), + total=len(self.trajectory_ids), + desc=f"Caching {key} frames", + ): + video_path = self.get_video_path(trajectory_id, key) + frames = get_all_frames( + video_path.as_posix(), + video_backend=self.video_backend, + video_backend_kwargs=self.video_backend_kwargs, + resize_size=img_resize, + ) + assert frames.ndim == 4, f"Expected 4D array, got {frames.shape} array" + assert frames.shape[3] == 3, f"Expected 3 channels, got {frames.shape[3]} channels" + + # Apply image cropping if enabled and the video key is base_view + # Note: crop_obs_camera functionality has been removed + + # assert ( + # frames.shape[0] == trajectory_length + # ), f"Expected {trajectory_length} frames, got {frames.shape[0]} frames" + all_frames.append(frames) + cached_frames[key] = np.concatenate(all_frames, axis=0) + print(f"{key}: {cached_frames[key].shape}") + self.cached_frames = cached_frames + self.start_indices = np.cumsum(self.trajectory_lengths) - self.trajectory_lengths + + def get_video(self, trajectory_id: int, key: str, base_index: int) -> np.ndarray: + step_indices = self.delta_indices[key] + base_index + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Ensure the indices are within the valid range + # This is equivalent to padding the video with extra frames at the beginning and end + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, self.trajectory_lengths[trajectory_index] - 1) + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Get the sub-key + key = key.replace("video.", "") + # Calculate the absolute indices + absolute_indices = self.start_indices[trajectory_index] + step_indices + return self.cached_frames[key][absolute_indices] + + def get_step_data(self, trajectory_id: int, base_index: int) -> dict: + """Get the RAW data for a single step. No transforms are applied. + + Args: + trajectory_id (str): The ID of the trajectory. + base_index (int): The base index of the step. + + Returns: + dict: The data for the step. + """ + data = {} + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + # Get the data for all modalities + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + data[key] = self.get_data_by_modality(trajectory_id, modality, key, base_index) + return data + + def set_transforms_metadata(self, metadata: DatasetMetadata): + """Set the metadata for the transforms. This is useful for transforms that need to know the metadata, such as the normalization values.""" + if self.img_resize is not None: + all_video_keys = [key for key in self.modality_keys["video"]] + for key in metadata.modalities.video: + if key in all_video_keys: + metadata.modalities.video[key].resolution = self.img_resize + super().set_transforms_metadata(metadata) + + +def safe_hash(input_tuple): + # keep 128 bits of the hash + tuple_string = repr(input_tuple).encode("utf-8") + sha256 = hashlib.sha256() + sha256.update(tuple_string) + + seed = int(sha256.hexdigest(), 16) + + return seed & 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF + + +class MixtureSpecElement(BaseModel): + dataset_path: list[Path] | Path = Field(..., description="The path to the dataset.") + dataset_weight: float = Field(..., description="The weight of the dataset in the mixture.") + distribute_weights: bool = Field( + default=False, + description="Whether to distribute the weights of the dataset across all the paths. If True, the weights will be evenly distributed across all the paths.", + ) + + +# Helper functions for dataset statistics + +def combine_modality_stats(modality_stats: dict) -> dict: + """ + Combine statistics from all sub-keys under a modality. + + Args: + modality_stats (dict): Statistics for a modality, containing multiple sub-keys. + Each sub-key contains DatasetStatisticalValues object. + + Returns: + dict: Combined statistics + """ + combined_stats = { + "mean": [], + "std": [], + "max": [], + "min": [], + "q01": [], + "q99": [] + } + + # Combine statistics in sub-key order + for subkey in modality_stats.keys(): + subkey_stats = modality_stats[subkey] # This is a DatasetStatisticalValues object + + # Convert DatasetStatisticalValues to dict-like access + for stat_name in ["mean", "std", "max", "min", "q01", "q99"]: + stat_value = getattr(subkey_stats, stat_name) + if isinstance(stat_value, (list, tuple)): + combined_stats[stat_name].extend(stat_value) + else: + # Handle NDArray case - convert to list + if hasattr(stat_value, 'tolist'): + combined_stats[stat_name].extend(stat_value.tolist()) + else: + combined_stats[stat_name].append(float(stat_value)) + + return combined_stats + +def _extract_action_normalization_modes(transforms) -> dict: + """Extract normalization modes for action keys from a ComposedModalityTransform. + + Returns: + dict: {subkey_without_action_prefix -> normalization_mode_str} + """ + modes = {} + for t in transforms.transforms: + if isinstance(t, StateActionTransform): + for key, mode in t.normalization_modes.items(): + if key.startswith("action."): + subkey = key[len("action."):] + modes[subkey] = mode + return modes + + +def generate_action_mask_for_used_keys( + action_modalities: dict, + used_action_keys_ordered, + normalization_modes: dict | None = None, +) -> list[bool]: + """Generate per-dimension mask for action statistics. + + A dimension gets ``mask=False`` only when its normalization mode is ``"binary"``. + This tells the inference code to skip continuous de-normalization for that dimension. + All other modes (q99, mean_std, min_max ...) produce ``mask=True``. + + Args: + action_modalities (dict): Configuration information for action modalities. + used_action_keys_ordered: Iterable of actually used action keys (no "action." prefix). + normalization_modes (dict | None): Mapping {subkey -> mode} (no "action." prefix). + If ``None``, all dimensions default to ``mask=True``. + + Returns: + list[bool]: Per-dimension mask values. + """ + mask = [] + + for subkey in used_action_keys_ordered: + if subkey in action_modalities: + subkey_config = action_modalities[subkey] + + # Get dimension count from shape + if hasattr(subkey_config, 'shape') and len(subkey_config.shape) > 0: + dim_count = subkey_config.shape[0] + else: + dim_count = 1 + + # mask=False only when the normalization mode is explicitly "binary" + is_binary = ( + normalization_modes.get(subkey) == "binary" + if normalization_modes is not None + else False + ) + + for _ in range(dim_count): + mask.append(not is_binary) + + return mask + +def get_used_modality_keys(modality_keys: dict) -> tuple[list, list]: + """Extract used action and state keys from modality configuration.""" + used_action_keys = [] + used_state_keys = [] + + # Extract action keys (remove "action." prefix) + for action_key in modality_keys.get("action", []): + if action_key.startswith("action."): + clean_key = action_key.replace("action.", "") + used_action_keys.append(clean_key) + + # Extract state keys (remove "state." prefix) + for state_key in modality_keys.get("state", []): + if state_key.startswith("state."): + clean_key = state_key.replace("state.", "") + used_state_keys.append(clean_key) + + return used_action_keys, used_state_keys + +class LeRobotMixtureDataset(Dataset): + """ + A mixture of multiple datasets. This class samples a single dataset based on the dataset weights and then calls the `__getitem__` method of the sampled dataset. + It is recommended to modify the single dataset class instead of this class. + """ + + def __init__( + self, + data_mixture: Sequence[tuple[LeRobotSingleDataset, float]], + mode: str, + balance_dataset_weights: bool = True, + balance_trajectory_weights: bool = True, + seed: int = 42, + metadata_config: dict = { + "percentile_mixing_method": "min_max", + }, + **kwargs, + ): + """ + Initialize the mixture dataset. + + Args: + data_mixture (list[tuple[LeRobotSingleDataset, float]]): Datasets and their corresponding weights. + mode (str): If "train", __getitem__ will return different samples every epoch; if "val" or "test", __getitem__ will return the same sample every epoch. + balance_dataset_weights (bool): If True, the weight of dataset will be multiplied by the total trajectory length of each dataset. + balance_trajectory_weights (bool): If True, sample trajectories within a dataset weighted by their length; otherwise, use equal weighting. + seed (int): Random seed for sampling. + """ + datasets: list[LeRobotSingleDataset] = [] + dataset_sampling_weights: list[float] = [] + for dataset, weight in data_mixture: + # Check if dataset is valid and has data + if len(dataset) == 0: + print(f"Warning: Skipping empty dataset {dataset.dataset_name}") + continue + datasets.append(dataset) + dataset_sampling_weights.append(weight) + + if len(datasets) == 0: + raise ValueError("No valid datasets found in the mixture. All datasets are empty.") + + self.datasets = datasets + self.balance_dataset_weights = balance_dataset_weights + self.balance_trajectory_weights = balance_trajectory_weights + self.seed = seed + self.mode = mode + self.data_cfg = kwargs["data_cfg"] if "data_cfg" in kwargs else None + self._traceable_sampling = bool( + self.data_cfg is not None and self.data_cfg.get("traceable_sampling", False) + ) + self.valid_anchor_manifest = None + self.temporal_proposal = None + + # Set properties for sampling + + # 1. Dataset lengths + self._dataset_lengths = np.array([len(dataset) for dataset in self.datasets]) + print(f"Dataset lengths: {self._dataset_lengths}") + self._getitem_count = 0 + # 2. Dataset sampling weights + self._dataset_sampling_weights = np.array(dataset_sampling_weights) + + if self.balance_dataset_weights: + self._dataset_sampling_weights *= self._dataset_lengths + + # Check for zero or negative weights before normalization + if np.any(self._dataset_sampling_weights <= 0): + print(f"Warning: Found zero or negative sampling weights: {self._dataset_sampling_weights}") + # Set minimum weight to prevent division issues + self._dataset_sampling_weights = np.maximum(self._dataset_sampling_weights, 1e-8) + + # Normalize weights + weights_sum = self._dataset_sampling_weights.sum() + if weights_sum == 0 or np.isnan(weights_sum): + print(f"Error: Invalid weights sum: {weights_sum}") + # Fallback to equal weights + self._dataset_sampling_weights = np.ones(len(self.datasets)) / len(self.datasets) + print(f"Fallback to equal weights") + else: + self._dataset_sampling_weights /= weights_sum + + # 3. Trajectory sampling weights + self._trajectory_sampling_weights: list[np.ndarray] = [] + for i, dataset in enumerate(self.datasets): + trajectory_sampling_weights = np.ones(len(dataset.trajectory_lengths)) + if self.balance_trajectory_weights: + trajectory_sampling_weights *= dataset.trajectory_lengths + + # Check for zero or negative weights before normalization + if np.any(trajectory_sampling_weights <= 0): + print(f"Warning: Dataset {i} has zero or negative trajectory weights") + trajectory_sampling_weights = np.maximum(trajectory_sampling_weights, 1e-8) + + # Normalize weights + weights_sum = trajectory_sampling_weights.sum() + if weights_sum == 0 or np.isnan(weights_sum): + print(f"Error: Dataset {i} has invalid trajectory weights sum: {weights_sum}") + # Fallback to equal weights + trajectory_sampling_weights = np.ones(len(dataset.trajectory_lengths)) / len(dataset.trajectory_lengths) + else: + trajectory_sampling_weights /= weights_sum + + self._trajectory_sampling_weights.append(trajectory_sampling_weights) + + if self._traceable_sampling: + self.valid_anchor_manifest = self._build_valid_anchor_manifest() + temporal_proposal_path = self.data_cfg.get("temporal_proposal_path") + if temporal_proposal_path: + self.temporal_proposal = TemporalProposal.load( + temporal_proposal_path, self.valid_anchor_manifest + ) + + # 4. Primary dataset indices + self._primary_dataset_indices = np.array(dataset_sampling_weights) == 1.0 + if not np.any(self._primary_dataset_indices): + print(f"Warning: No dataset with weight 1.0 found. Original weights: {dataset_sampling_weights}") + # Fallback: use the dataset(s) with maximum weight as primary + max_weight = max(dataset_sampling_weights) + self._primary_dataset_indices = np.array(dataset_sampling_weights) == max_weight + print(f"Using datasets with maximum weight {max_weight} as primary: {self._primary_dataset_indices}") + + if not np.any(self._primary_dataset_indices): + # This should never happen, but just in case + print("Error: Still no primary dataset found. Using first dataset as primary.") + self._primary_dataset_indices = np.zeros(len(self.datasets), dtype=bool) + self._primary_dataset_indices[0] = True + + # Set the epoch and sample the first epoch + self.set_epoch(0) + + self.update_metadata(metadata_config) + + @property + def dataset_lengths(self) -> np.ndarray: + """The lengths of each dataset.""" + return self._dataset_lengths + + @property + def dataset_sampling_weights(self) -> np.ndarray: + """The sampling weights for each dataset.""" + return self._dataset_sampling_weights + + @property + def trajectory_sampling_weights(self) -> list[np.ndarray]: + """The sampling weights for each trajectory in each dataset.""" + return self._trajectory_sampling_weights + + @property + def primary_dataset_indices(self) -> np.ndarray: + """The indices of the primary datasets.""" + return self._primary_dataset_indices + + def __str__(self) -> str: + dataset_descriptions = [] + for dataset, weight in zip(self.datasets, self.dataset_sampling_weights): + dataset_description = { + "Dataset": str(dataset), + "Sampling weight": float(weight), + } + dataset_descriptions.append(dataset_description) + return json.dumps({"Mixture dataset": dataset_descriptions}, indent=2) + + def set_epoch(self, epoch: int): + """Set the epoch for the dataset. + + Args: + epoch (int): The epoch to set. + """ + self.epoch = epoch + # self.sampled_steps = self.sample_epoch() + + def _build_valid_anchor_manifest(self) -> ValidAnchorManifest: + """Freeze the readable episode/anchor support used by probability-aware training.""" + def serialize_transforms(dataset) -> list[dict]: + serialized = [] + for transform in dataset.transforms.transforms: + try: + config = BaseModel.model_dump(transform, mode="json") + except Exception: + config = BaseModel.model_dump(transform, mode="python") + config = json.loads(json.dumps(config, default=str, sort_keys=True)) + serialized.append( + { + "class": f"{type(transform).__module__}.{type(transform).__qualname__}", + "config": config, + } + ) + return serialized + + all_dataset_episodes = [] + excluded_episodes = [] + for dataset_index, dataset in enumerate(self.datasets): + anchors_by_trajectory = defaultdict(list) + for trajectory_id, base_index in dataset.all_steps: + trajectory_key = trajectory_id.item() if isinstance(trajectory_id, np.generic) else trajectory_id + anchors_by_trajectory[trajectory_key].append(int(base_index)) + + valid_episodes = [] + total_videos = int(dataset.lerobot_info_meta.get("total_videos", 0)) + for trajectory_index, trajectory_id in enumerate(dataset.trajectory_ids): + trajectory_key = trajectory_id.item() if isinstance(trajectory_id, np.generic) else trajectory_id + anchors = anchors_by_trajectory.get(trajectory_key, []) + if not anchors: + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "no_valid_anchors", + } + ) + continue + + if dataset._lerobot_version == "v2.0": + parquet_path = dataset.dataset_path / dataset.data_path_pattern.format( + episode_chunk=dataset.get_episode_chunk(trajectory_key), + episode_index=trajectory_key, + ) + else: + episode_metadata = dataset.trajectory_ids_to_metadata[trajectory_key] + parquet_path = dataset.dataset_path / dataset.data_path_pattern.format( + chunk_index=episode_metadata["data/chunk_index"], + file_index=episode_metadata["data/file_index"], + ) + if not parquet_path.exists(): + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "missing_parquet", + "path": str(parquet_path), + } + ) + continue + try: + import pyarrow.parquet as parquet + + parquet.ParquetFile(parquet_path) + except Exception as exc: + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "unreadable_parquet", + "path": str(parquet_path), + "error": f"{type(exc).__name__}: {exc}", + } + ) + continue + + if total_videos > 0: + video_paths = [ + dataset.get_video_path(trajectory_key, key.replace("video.", "")) + for key in dataset.modality_keys["video"] + ] + if not all(path.exists() for path in video_paths): + excluded_episodes.append( + { + "dataset": dataset.dataset_name, + "trajectory_id": trajectory_key, + "reason": "missing_video", + "paths": [str(path) for path in video_paths if not path.exists()], + } + ) + continue + + valid_episodes.append( + ( + trajectory_key, + tuple(sorted(set(anchors))), + float(self.trajectory_sampling_weights[dataset_index][trajectory_index]), + ) + ) + + if not valid_episodes: + raise RuntimeError( + f"Traceable sampling found no valid episodes in dataset {dataset.dataset_name}" + ) + all_dataset_episodes.append(valid_episodes) + + data_config = {} + if self.data_cfg is not None: + manifest_config_keys = ( + "data_mix", + "delete_pause_frame", + "lerobot_version", + "action_type", + "action_mode", + "action_mode_apply_keys", + "action_mode_state_map", + "include_state", + "video_backend", + "precomputed_statistics_path", + "precomputed_statistics_key_map", + "precomputed_statistics_filename", + ) + try: + from omegaconf import OmegaConf + + resolved_data_config = OmegaConf.to_container(self.data_cfg, resolve=True) + except (TypeError, ValueError): + resolved_data_config = dict(self.data_cfg) + data_config = { + key: resolved_data_config.get(key) + for key in manifest_config_keys + if key in resolved_data_config + } + + dataset_provenance = None + provenance_path = self.data_cfg.get("dataset_provenance_path") if self.data_cfg else None + if provenance_path: + with open(provenance_path, "r", encoding="utf-8") as provenance_file: + dataset_provenance = json.load(provenance_file) + + content_fingerprint_path = ( + self.data_cfg.get("dataset_content_fingerprint_path") if self.data_cfg else None + ) + if content_fingerprint_path: + with open(content_fingerprint_path, "r", encoding="utf-8") as fingerprint_file: + content_fingerprints = json.load(fingerprint_file) + else: + content_fingerprints = {"datasets": { + dataset.dataset_name: compute_dataset_content_fingerprint(dataset.dataset_path) + for dataset in self.datasets + }} + dataset_identities = [] + for dataset in self.datasets: + try: + content_fingerprint = content_fingerprints["datasets"][dataset.dataset_name] + except KeyError as exc: + raise ValueError( + f"No content fingerprint found for dataset {dataset.dataset_name}" + ) from exc + validate_dataset_content_inventory(dataset.dataset_path, content_fingerprint) + dataset_identities.append( + { + "name": dataset.dataset_name, + "info": dataset.lerobot_info_meta, + "metadata": dataset.metadata.model_dump(mode="json"), + "modality": dataset._lerobot_modality_meta.model_dump(mode="json"), + "transforms": serialize_transforms(dataset), + "content": { + "file_count": content_fingerprint["file_count"], + "total_bytes": content_fingerprint["total_bytes"], + "files": content_fingerprint["files"], + }, + "delta_indices": { + key: value.tolist() for key, value in dataset.delta_indices.items() + }, + } + ) + + identity = { + "data_config": data_config, + "balance_dataset_weights": self.balance_dataset_weights, + "balance_trajectory_weights": self.balance_trajectory_weights, + "excluded_episodes": excluded_episodes, + "dataset_provenance": dataset_provenance, + "datasets": dataset_identities, + } + identity = json.loads(json.dumps(identity, sort_keys=True, default=str)) + current_manifest = ValidAnchorManifest.from_hierarchy( + dataset_names=[dataset.dataset_name for dataset in self.datasets], + dataset_weights=self.dataset_sampling_weights, + episodes=all_dataset_episodes, + identity=identity, + ) + manifest_path = self.data_cfg.get("valid_anchor_manifest_path") if self.data_cfg else None + if manifest_path: + frozen_manifest = ValidAnchorManifest.load(manifest_path) + frozen_support = frozen_manifest.to_dict(include_hash=False)["datasets"] + current_support = current_manifest.to_dict(include_hash=False)["datasets"] + if frozen_support != current_support: + raise RuntimeError( + "Frozen valid-anchor manifest changes dataset/episode/anchor support" + ) + frozen_identity = json.loads(json.dumps(frozen_manifest.identity)) + current_identity = json.loads(json.dumps(current_manifest.identity)) + for identity in (frozen_identity, current_identity): + provenance = identity.get("dataset_provenance") + if isinstance(provenance, dict) and "payload" in provenance: + identity["dataset_provenance"] = provenance["payload"] + identity.pop("content_fingerprint_artifact_sha256", None) + for dataset_identity in identity.get("datasets", []): + content = dataset_identity.get("content", {}) + content.pop("sha256", None) + data_config = identity.get("data_config", {}) + statistics_path = data_config.get("precomputed_statistics_path") + if statistics_path: + data_config["precomputed_statistics_path"] = Path(statistics_path).name + if frozen_identity != current_identity: + raise RuntimeError( + "Frozen valid-anchor manifest changes the semantic data configuration" + ) + return frozen_manifest + return current_manifest + + def sample_trace(self, index: int | LogicalDraw): + if self.valid_anchor_manifest is None: + raise RuntimeError("sample_trace requires traceable_sampling=True") + draw = index if isinstance(index, LogicalDraw) else LogicalDraw(self.epoch, int(index)) + if self.temporal_proposal is not None: + return self.temporal_proposal.sample(self.valid_anchor_manifest, self.seed, draw) + return self.valid_anchor_manifest.sample(self.seed, draw) + + def sample_step(self, index: int) -> tuple[LeRobotSingleDataset, int, int]: + """Sample a single step from the dataset.""" + # return self.sampled_steps[index] + + if self.valid_anchor_manifest is not None: + trace = self.sample_trace(index) + return self.datasets[trace.dataset_index], trace.trajectory_id, trace.base_index + + # Set seed + seed = index if self.mode != "train" else safe_hash((self.epoch, index, self.seed)) + rng = np.random.default_rng(seed) + + # Sample dataset + dataset_index = rng.choice(len(self.datasets), p=self.dataset_sampling_weights) + dataset = self.datasets[dataset_index] + + # Sample trajectory + trajectory_index = rng.choice( + len(dataset.trajectory_ids), p=self.trajectory_sampling_weights[dataset_index] + ) + trajectory_id = dataset.trajectory_ids[trajectory_index] + + # Sample step + base_index = rng.choice(dataset.trajectory_lengths[trajectory_index]) + return dataset, trajectory_id, base_index + + + + def __getitem__(self, index: int) -> dict: + """Get the data for a single trajectory and start index. + + Args: + index (int): The index of the trajectory to get. + + Returns: + dict: The data for the trajectory and start index. + """ + self._getitem_count += 1 + if self._getitem_count % 1000 == 0: + gc.collect() + + if self.valid_anchor_manifest is not None: + trace = self.sample_trace(index) + dataset = self.datasets[trace.dataset_index] + try: + raw_data = dataset.get_step_data(trace.trajectory_id, trace.base_index) + transform_rng_seed = ( + self.data_cfg.get("transform_rng_seed") + if self.data_cfg is not None + else None + ) + logical_transform_seed = None + if transform_rng_seed is not None: + payload = ( + f"{int(transform_rng_seed)}:{trace.epoch}:" + f"{trace.global_draw_id}" + ).encode("utf-8") + logical_transform_seed = int.from_bytes( + hashlib.sha256(payload).digest()[:8], + "little", + signed=False, + ) + with _isolated_transform_rng(logical_transform_seed): + data = dataset.transforms(raw_data) + sample = dataset._pack_sample(data) + except Exception as exc: + raise RuntimeError( + "A frozen valid-anchor manifest entry failed at runtime; refusing to " + f"change the sampling distribution. trace={trace.to_dict()}" + ) from exc + sample["sampling_trace"] = trace.to_dict() + return sample + + max_retries = 10 + last_exception = None + + for attempt in range(max_retries): + try: + sample_tries = 0 + max_sample_tries = 200 + while True: + sample_tries += 1 + if sample_tries > max_sample_tries: + raise RuntimeError( + f"Unable to sample a valid item after {max_sample_tries} attempts. " + f"dataset={self.datasets[0].dataset_name if len(self.datasets)>0 else 'unknown'}" + ) + + dataset, trajectory_id, step = self.sample_step(index) + # If dataset has no physical videos (e.g., image frames in parquet + # for VLA-Arena), sampling does not require mp4 files. + total_videos = int(dataset.lerobot_info_meta.get("total_videos", 0)) + if total_videos == 0: + break + + key = dataset.modality_keys["video"][0].replace("video.", "") + video_path = dataset.get_video_path(trajectory_id, key) + if os.path.exists(video_path): + break + index = random.randint(0, len(self) - 1) + + raw_data = dataset.get_step_data(trajectory_id, step) + data = dataset.transforms(raw_data) + sample = dataset._pack_sample(data) + + return sample + + except Exception as e: + last_exception = e + if attempt < max_retries - 1: + # Log the error but continue trying + print(f"Attempt {attempt + 1}/{max_retries} failed for index {index}: {e}") + print(f"Retrying with new sample...") + # For retry, we can use a slightly different index to get a new sample + # This helps avoid getting stuck on the same problematic sample + index = random.randint(0, len(self) - 1) + else: + # All retries exhausted + print(f"All {max_retries} attempts failed for index {index}") + print(f"Last error: {last_exception}") + # Return a dummy sample or re-raise the exception + raise last_exception + + def __len__(self) -> int: + """Get the length of a single epoch in the mixture. + + Returns: + int: The length of a single epoch in the mixture. + """ + # Check for potential issues + if len(self.datasets) == 0: + return 0 + + # Check if any dataset lengths are 0 or NaN + if np.any(self.dataset_lengths == 0) or np.any(np.isnan(self.dataset_lengths)): + print(f"Warning: Found zero or NaN dataset lengths: {self.dataset_lengths}") + # Filter out zero/NaN length datasets + valid_indices = (self.dataset_lengths > 0) & (~np.isnan(self.dataset_lengths)) + if not np.any(valid_indices): + print("Error: All datasets have zero or NaN length") + return 0 + else: + valid_indices = np.ones(len(self.datasets), dtype=bool) + + # Check if any sampling weights are 0 or NaN + if np.any(self.dataset_sampling_weights == 0) or np.any(np.isnan(self.dataset_sampling_weights)): + print(f"Warning: Found zero or NaN sampling weights: {self.dataset_sampling_weights}") + # Use only valid weights + valid_weights = (self.dataset_sampling_weights > 0) & (~np.isnan(self.dataset_sampling_weights)) + valid_indices = valid_indices & valid_weights + if not np.any(valid_indices): + print("Error: All sampling weights are zero or NaN") + return 0 + + # Check primary dataset indices + primary_and_valid = self.primary_dataset_indices & valid_indices + if not np.any(primary_and_valid): + print(f"Warning: No valid primary datasets found. Primary indices: {self.primary_dataset_indices}, Valid indices: {valid_indices}") + # Fallback: use the largest valid dataset + if np.any(valid_indices): + max_length = self.dataset_lengths[valid_indices].max() + print(f"Fallback: Using maximum dataset length: {max_length}") + return int(max_length) + else: + return 0 + + # Calculate the ratio and get max + ratios = (self.dataset_lengths / self.dataset_sampling_weights)[primary_and_valid] + + # Check for NaN or inf in ratios + if np.any(np.isnan(ratios)) or np.any(np.isinf(ratios)): + print(f"Warning: Found NaN or inf in ratios: {ratios}") + print(f"Dataset lengths: {self.dataset_lengths[primary_and_valid]}") + print(f"Sampling weights: {self.dataset_sampling_weights[primary_and_valid]}") + # Filter out invalid ratios + valid_ratios = ratios[~np.isnan(ratios) & ~np.isinf(ratios)] + if len(valid_ratios) == 0: + print("Error: All ratios are NaN or inf") + return 0 + max_ratio = valid_ratios.max() + else: + max_ratio = ratios.max() + + result = int(max_ratio) + if result == 0: + print(f"Warning: Dataset mixture length is 0") + return result + + @staticmethod + def compute_overall_statistics( + per_task_stats: list[dict[str, dict[str, list[float] | np.ndarray]]], + dataset_sampling_weights: list[float] | np.ndarray, + percentile_mixing_method: str = "weighted_average", + ) -> dict[str, dict[str, list[float]]]: + """ + Computes overall statistics from per-task statistics using dataset sample weights. + + Args: + per_task_stats: List of per-task statistics. + Example format of one element in the per-task statistics list: + { + "state.gripper": { + "min": [...], + "max": [...], + "mean": [...], + "std": [...], + "q01": [...], + "q99": [...], + }, + ... + } + dataset_sampling_weights: List of sample weights for each task. + percentile_mixing_method: The method to mix the percentiles, either "weighted_average" or "weighted_std". + + Returns: + A dict of overall statistics per modality. + """ + # Normalize the sample weights to sum to 1 + dataset_sampling_weights = np.array(dataset_sampling_weights) + normalized_weights = dataset_sampling_weights / dataset_sampling_weights.sum() + + # Initialize overall statistics dict + overall_stats: dict[str, dict[str, list[float]]] = {} + + # Get the list of modality keys + modality_keys = per_task_stats[0].keys() + + for modality in modality_keys: + # Number of dimensions (assuming consistent across tasks) + num_dims = len(per_task_stats[0][modality]["mean"]) + + # Initialize accumulators for means and variances + weighted_means = np.zeros(num_dims) + weighted_squares = np.zeros(num_dims) + + # Collect min, max, q01, q99 from all tasks + min_list = [] + max_list = [] + q01_list = [] + q99_list = [] + + for task_idx, task_stats in enumerate(per_task_stats): + w_i = normalized_weights[task_idx] + stats = task_stats[modality] + means = np.array(stats["mean"]) + stds = np.array(stats["std"]) + + # Update weighted sums for mean and variance + weighted_means += w_i * means + weighted_squares += w_i * (stds**2 + means**2) + + # Collect min, max, q01, q99 + min_list.append(stats["min"]) + max_list.append(stats["max"]) + q01_list.append(stats["q01"]) + q99_list.append(stats["q99"]) + + # Compute overall mean + overall_mean = weighted_means.tolist() + + # Compute overall variance and std deviation + overall_variance = weighted_squares - weighted_means**2 + overall_std = np.sqrt(overall_variance).tolist() + + # Compute overall min and max per dimension + overall_min = np.min(np.array(min_list), axis=0).tolist() + overall_max = np.max(np.array(max_list), axis=0).tolist() + + # Compute overall q01 and q99 per dimension + # Use weighted average of per-task quantiles + q01_array = np.array(q01_list) + q99_array = np.array(q99_list) + if percentile_mixing_method == "weighted_average": + weighted_q01 = np.average(q01_array, axis=0, weights=normalized_weights).tolist() + weighted_q99 = np.average(q99_array, axis=0, weights=normalized_weights).tolist() + # std_q01 = np.std(q01_array, axis=0).tolist() + # std_q99 = np.std(q99_array, axis=0).tolist() + # print(modality) + # print(f"{std_q01=}, {std_q99=}") + # print(f"{weighted_q01=}, {weighted_q99=}") + elif percentile_mixing_method == "min_max": + weighted_q01 = np.min(q01_array, axis=0).tolist() + weighted_q99 = np.max(q99_array, axis=0).tolist() + else: + raise ValueError(f"Invalid percentile mixing method: {percentile_mixing_method}") + + # Store the overall statistics for the modality + overall_stats[modality] = { + "min": overall_min, + "max": overall_max, + "mean": overall_mean, + "std": overall_std, + "q01": weighted_q01, + "q99": weighted_q99, + } + + return overall_stats + + @staticmethod + def merge_metadata( + metadatas: list[DatasetMetadata], + dataset_sampling_weights: list[float], + percentile_mixing_method: str, + ) -> DatasetMetadata: + """Merge multiple metadata into one.""" + # Convert to dicts + metadata_dicts = [metadata.model_dump(mode="json") for metadata in metadatas] + # Create a new metadata dict + merged_metadata = {} + + # Check all metadata have the same embodiment tag + assert all( + metadata.embodiment_tag == metadatas[0].embodiment_tag for metadata in metadatas + ), "All metadata must have the same embodiment tag" + merged_metadata["embodiment_tag"] = metadatas[0].embodiment_tag + + # Merge the dataset statistics + dataset_statistics = {} + dataset_statistics["state"] = LeRobotMixtureDataset.compute_overall_statistics( + per_task_stats=[m["statistics"]["state"] for m in metadata_dicts], + dataset_sampling_weights=dataset_sampling_weights, + percentile_mixing_method=percentile_mixing_method, + ) + dataset_statistics["action"] = LeRobotMixtureDataset.compute_overall_statistics( + per_task_stats=[m["statistics"]["action"] for m in metadata_dicts], + dataset_sampling_weights=dataset_sampling_weights, + percentile_mixing_method=percentile_mixing_method, + ) + merged_metadata["statistics"] = dataset_statistics + + # Merge the modality configs + modality_configs = defaultdict(set) + for metadata in metadata_dicts: + for modality, configs in metadata["modalities"].items(): + modality_configs[modality].add(json.dumps(configs)) + merged_metadata["modalities"] = {} + for modality, configs in modality_configs.items(): + # Check that all modality configs correspond to the same tag matches + assert ( + len(configs) == 1 + ), f"Multiple modality configs for modality {modality}: {list(configs)}" + merged_metadata["modalities"][modality] = json.loads(configs.pop()) + + return DatasetMetadata.model_validate(merged_metadata) + + def update_metadata(self, metadata_config: dict, cached_statistics_path: Path | str | None = None) -> None: + """ + Merge multiple metadatas into one and set the transforms with the merged metadata. + + Args: + metadata_config (dict): Configuration for the metadata. + "percentile_mixing_method": The method to mix the percentiles, either "weighted_average" or "min_max". + weighted_average: Use the weighted average of the percentiles using the weight used in sampling the datasets. + min_max: Use the min of the 1st percentile and max of the 99th percentile. + """ + # If cached path is provided, try to load and apply + if cached_statistics_path is not None: + try: + cached_stats = self.load_merged_statistics(cached_statistics_path) + self.apply_cached_statistics(cached_stats) + return + except (FileNotFoundError, KeyError, ValidationError) as e: + print(f"Failed to load cached statistics: {e}") + print("Falling back to computing statistics from scratch...") + + self.tag = EmbodimentTag.NEW_EMBODIMENT.value + self.merged_metadata: dict[str, DatasetMetadata] = {} + # Group metadata by tag + all_metadatas: dict[str, list[DatasetMetadata]] = {} + for dataset in self.datasets: + if dataset.tag not in all_metadatas: + all_metadatas[dataset.tag] = [] + all_metadatas[dataset.tag].append(dataset.metadata) + for tag, metadatas in all_metadatas.items(): + self.merged_metadata[tag] = self.merge_metadata( + metadatas=metadatas, + dataset_sampling_weights=self.dataset_sampling_weights.tolist(), + percentile_mixing_method=metadata_config["percentile_mixing_method"], + ) + for dataset in self.datasets: + dataset.set_transforms_metadata(self.merged_metadata[dataset.tag]) + + def save_dataset_statistics(self, save_path: Path | str, format: str = "json") -> None: + """ + Save merged dataset statistics to specified path in the required format. + Only includes statistics for keys that are actually used in the datasets. + Gripper-related keys will be placed at the end. + + Args: + save_path (Path | str): Path to save the statistics file + format (str): Save format, currently only supports "json" + """ + save_path = Path(save_path) + save_path.parent.mkdir(parents=True, exist_ok=True) + + # Build the data structure to save + statistics_data = {} + + # Collect actually used keys from all datasets + all_used_action_keys = [] + all_used_state_keys = [] + + for dataset in self.datasets: + used_action_keys, used_state_keys = get_used_modality_keys(dataset.modality_keys) + for used_action_key in used_action_keys: + if used_action_key not in all_used_action_keys: + all_used_action_keys.append(used_action_key) + for used_state_key in used_state_keys: + if used_state_key not in all_used_state_keys: + all_used_state_keys.append(used_state_key) + + # Organize statistics by tag + for tag, merged_metadata in self.merged_metadata.items(): + tag_stats = {} + + # Process action statistics + if hasattr(merged_metadata.statistics, 'action') and merged_metadata.statistics.action: + action_stats = merged_metadata.statistics.action + + # Filter and reorder keys - iterate in all_used_action_keys order + non_gripper_keys = [] + gripper_keys = [] + + for key in all_used_action_keys: + if key in action_stats: + non_gripper_keys.append(key) + + reordered_keys = non_gripper_keys + gripper_keys + + filtered_action_stats = {} + for key in reordered_keys: + filtered_action_stats[key] = action_stats[key] + + if filtered_action_stats: + combined_action_stats = combine_modality_stats(filtered_action_stats) + + # Collect action normalization modes from datasets of this tag. + # "binary" takes precedence: if any dataset marks a key as binary, use binary. + _action_norm_modes: dict = {} + for _ds in self.datasets: + if _ds.tag == tag: + for _k, _m in _extract_action_normalization_modes(_ds.transforms).items(): + if _k not in _action_norm_modes or _m == "binary": + _action_norm_modes[_k] = _m + mask = generate_action_mask_for_used_keys( + merged_metadata.modalities.action, filtered_action_stats.keys(), + normalization_modes=_action_norm_modes, + ) + combined_action_stats["mask"] = mask + + tag_stats["action"] = combined_action_stats + + # Process state statistics + if hasattr(merged_metadata.statistics, 'state') and merged_metadata.statistics.state: + state_stats = merged_metadata.statistics.state + + # Filter and reorder keys - iterate in all_used_state_keys order + # Filter and reorder keys - iterate in all_used_state_keys order + non_gripper_keys = [] + gripper_keys = [] + + for key in all_used_state_keys: + if key in state_stats: + non_gripper_keys.append(key) + + reordered_keys = non_gripper_keys + gripper_keys + + filtered_state_stats = {} + for key in reordered_keys: + filtered_state_stats[key] = state_stats[key] + + if filtered_state_stats: + combined_state_stats = combine_modality_stats(filtered_state_stats) + tag_stats["state"] = combined_state_stats + + # Add dataset counts + tag_stats.update(self._get_dataset_counts(tag)) + + statistics_data[tag] = tag_stats + + # Save file + if format.lower() == "json": + if not str(save_path).endswith('.json'): + save_path = save_path.with_suffix('.json') + with open(save_path, 'w', encoding='utf-8') as f: + json.dump(statistics_data, f, indent=2, ensure_ascii=False) + else: + raise ValueError(f"Unsupported format: {format}. Currently only 'json' is supported.") + + print(f"Merged dataset statistics saved to: {save_path}") + print(f"Used action keys (reordered): {list(all_used_action_keys)}") + print(f"Used state keys (reordered): {list(all_used_state_keys)}") + + + def _combine_modality_stats(self, modality_stats: dict) -> dict: + """Backward compatibility wrapper.""" + return combine_modality_stats(modality_stats) + + def _generate_action_mask_for_used_keys(self, action_modalities: dict, used_action_keys_ordered) -> list[bool]: + """Backward compatibility wrapper.""" + return generate_action_mask_for_used_keys(action_modalities, used_action_keys_ordered) + + def _get_dataset_counts(self, tag: str) -> dict: + """ + Get dataset count information for specified tag. + + Args: + tag (str): embodiment tag + + Returns: + dict: Dictionary containing num_transitions and num_trajectories + """ + num_transitions = 0 + num_trajectories = 0 + + # Count dataset information belonging to this tag + for dataset in self.datasets: + if dataset.tag == tag: + num_transitions += len(dataset) + num_trajectories += len(dataset.trajectory_ids) + + return { + "num_transitions": num_transitions, + "num_trajectories": num_trajectories + } + + @classmethod + def load_merged_statistics(cls, load_path: Path | str) -> dict: + """ + Load merged dataset statistics from file. + + Args: + load_path (Path | str): Path to the statistics file + + Returns: + dict: Dictionary containing merged statistics + """ + load_path = Path(load_path) + if not load_path.exists(): + raise FileNotFoundError(f"Statistics file not found: {load_path}") + + if load_path.suffix.lower() == '.json': + with open(load_path, 'r', encoding='utf-8') as f: + return json.load(f) + elif load_path.suffix.lower() == '.pkl': + import pickle + with open(load_path, 'rb') as f: + return pickle.load(f) + else: + raise ValueError(f"Unsupported file format: {load_path.suffix}") + + def apply_cached_statistics(self, cached_statistics: dict) -> None: + """ + Apply cached statistics to avoid recomputation. + + Args: + cached_statistics (dict): Statistics loaded from file + """ + # Validate that cached statistics match current datasets + if "metadata" in cached_statistics: + cached_dataset_names = set(cached_statistics["metadata"]["dataset_names"]) + current_dataset_names = set(dataset.dataset_name for dataset in self.datasets) + + if cached_dataset_names != current_dataset_names: + print("Warning: Cached statistics dataset names don't match current datasets.") + print(f"Cached: {cached_dataset_names}") + print(f"Current: {current_dataset_names}") + return + + # Apply cached statistics + self.merged_metadata = {} + for tag, stats_data in cached_statistics.items(): + if tag == "metadata": # Skip metadata field + continue + + # Convert back to DatasetMetadata format + metadata_dict = { + "embodiment_tag": tag, + "statistics": { + "action": {}, + "state": {} + }, + "modalities": {} + } + + # Convert action statistics back + if "action" in stats_data: + action_data = stats_data["action"] + # This is simplified - you may need to split back to sub-keys + metadata_dict["statistics"]["action"] = action_data + + # Convert state statistics back + if "state" in stats_data: + state_data = stats_data["state"] + metadata_dict["statistics"]["state"] = state_data + + self.merged_metadata[tag] = DatasetMetadata.model_validate(metadata_dict) + + # Update transforms metadata for each dataset + for dataset in self.datasets: + if dataset.tag in self.merged_metadata: + dataset.set_transforms_metadata(self.merged_metadata[dataset.tag]) + + print(f"Applied cached statistics for {len(self.merged_metadata)} embodiment tags.") diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py new file mode 100644 index 0000000000000000000000000000000000000000..b4cb43804eb0fd9c533df2addb27e0b57f62b340 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/embodiment_tags.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 enum import Enum + + +class EmbodimentTag(Enum): + GR1 = "gr1" + """ + The GR1 dataset. + """ + + OXE_DROID = "oxe_droid" + """ + The OxE Droid dataset. + """ + + OXE_BRIDGE = "oxe_bridge" + """ + The OxE Bridge dataset. + """ + + OXE_RT1 = "oxe_rt1" + """ + The OxE RT-1 dataset. + """ + + AGIBOT_GENIE1 = "agibot_genie1" + """ + The AgiBot Genie-1 with gripper dataset. + """ + + NEW_EMBODIMENT = "new_embodiment" + """ + Any new embodiment for finetuning. + """ + + FRANKA = 'franka' + """ + The Franka Emika Panda robot. + """ + + ALOHA = 'aloha' + """ + The ALOHA bimanual robot (RoboChallenge dual-arm). + """ + + UR5 = 'ur5' + """ The UR5 single-arm robot. + """ + + ARX5 = 'arx5' + """ The ARX5 single-arm robot. + """ + + DOS_W1 = 'dos-w1' + """ The DOS-W1 single-arm robot. + """ + +# Embodiment tag string: to projector index in the Action Expert Module +EMBODIMENT_TAG_MAPPING = { + EmbodimentTag.NEW_EMBODIMENT.value: 31, + EmbodimentTag.OXE_DROID.value: 17, + EmbodimentTag.OXE_BRIDGE.value: 18, + EmbodimentTag.OXE_RT1.value: 19, + EmbodimentTag.AGIBOT_GENIE1.value: 26, + EmbodimentTag.GR1.value: 24, + EmbodimentTag.FRANKA.value: 25, + EmbodimentTag.ALOHA.value: 7, + EmbodimentTag.UR5.value: 8, + EmbodimentTag.ARX5.value: 9, + EmbodimentTag.DOS_W1.value: 10, +} + +# Robot type to embodiment tag mapping +ROBOT_TYPE_TO_EMBODIMENT_TAG = { + "libero_franka": EmbodimentTag.FRANKA, + "oxe_droid": EmbodimentTag.OXE_DROID, + "oxe_bridge": EmbodimentTag.OXE_BRIDGE, + "oxe_rt1": EmbodimentTag.OXE_RT1, + "demo_sim_franka_delta_joints": EmbodimentTag.FRANKA, + "custom_robot_config": EmbodimentTag.NEW_EMBODIMENT, + "fourier_gr1_arms_waist": EmbodimentTag.GR1, +} diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/mixtures.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/mixtures.py new file mode 100644 index 0000000000000000000000000000000000000000..40957c507912fd204d8f9e3191ab807fcc7ecc93 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/mixtures.py @@ -0,0 +1,376 @@ +""" +mixtures.py + +Defines a registry of dataset mixtures and weights for the Open-X Embodiment Datasets. Each dataset is associated with +a float "sampling weight" +""" + +from typing import Dict, List, Tuple + + +# Dataset mixture name mapped to a list of tuples containing: +## {nakename: [(data_name, sampling_weight, robot_type)] } +DATASET_NAMED_MIXTURES = { + + "custom_dataset": [ + ("custom_dataset_name", 1.0, "custom_robot_config"), + ], + "custom_dataset_2": [ + ("custom_dataset_name_1", 1.0, "custom_robot_config"), + ("custom_dataset_name_2", 1.0, "custom_robot_config"), + ], + + "libero_all": [ + ("libero_object_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ("libero_goal_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ("libero_spatial_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ("libero_10_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + # ("libero_90_no_noops_lerobot", 1.0, "libero_franka"), + ], + "libero_goal": [ + ("libero_goal_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + ], + "bridge": [ + ("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"), + ], + "bridge_rt_1": [ + ("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"), + ("fractal20220817_data_0.1.0_lerobot", 1.0, "oxe_rt1"), + ], + + "demo_sim_pick_place": [ + ("sim_pick_place", 1.0, "demo_sim_franka_delta_joints"), + ], + + "custom_dataset": [ + ("custom_dataset_name", 1.0, "custom_robot_config"), + ], + "custom_dataset_2": [ + ("custom_dataset_name_1", 1.0, "custom_robot_config"), + ("custom_dataset_name_2", 1.0, "custom_robot_config"), + ], + + "fourier_gr1_unified_1000": [ + ("gr1_unified.PnPBottleToCabinetClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPCanToDrawerClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPCupToDrawerClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPMilkToMicrowaveClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPPotatoToMicrowaveClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PnPWineToCabinetClose_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToBasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToPanSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToPotSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromCuttingboardToTieredbasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToBasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToBowlSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToPlateSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlacematToTieredshelfSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToBowlSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToPanSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromPlateToPlateSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToPlateSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToPotSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToTieredbasketSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ("gr1_unified.PosttrainPnPNovelFromTrayToTieredshelfSplitA_GR1ArmsAndWaistFourierHands_1000", 1.0, "fourier_gr1_arms_waist"), + ], + + "BEHAVIOR_challenge": [ + ("BEHAVIOR_challenge", 1.0, "R1Pro"), + ], + + + "SO101_pick": [ + ("pick_dataset_name", 1.0, "SO101"), + ], + + "arx_x5": [ + ("arx_x5", 1.0, "arx_x5"), + ], + + "robotwin_all": [ + ("Clean/adjust_bottle", 1.0, "robotwin"), + ("Clean/beat_block_hammer", 1.0, "robotwin"), + ("Clean/blocks_ranking_rgb", 1.0, "robotwin"), + ("Clean/blocks_ranking_size", 1.0, "robotwin"), + ("Clean/click_alarmclock", 1.0, "robotwin"), + ("Clean/click_bell", 1.0, "robotwin"), + ("Clean/dump_bin_bigbin", 1.0, "robotwin"), + ("Clean/grab_roller", 1.0, "robotwin"), + ("Clean/handover_block", 1.0, "robotwin"), + ("Clean/handover_mic", 1.0, "robotwin"), + ("Clean/hanging_mug", 1.0, "robotwin"), + ("Clean/lift_pot", 1.0, "robotwin"), + ("Clean/move_can_pot", 1.0, "robotwin"), + ("Clean/move_pillbottle_pad", 1.0, "robotwin"), + ("Clean/move_playingcard_away", 1.0, "robotwin"), + ("Clean/move_stapler_pad", 1.0, "robotwin"), + ("Clean/open_laptop", 1.0, "robotwin"), + ("Clean/open_microwave", 1.0, "robotwin"), + ("Clean/pick_diverse_bottles", 1.0, "robotwin"), + ("Clean/pick_dual_bottles", 1.0, "robotwin"), + ("Clean/place_a2b_left", 1.0, "robotwin"), + ("Clean/place_a2b_right", 1.0, "robotwin"), + ("Clean/place_bread_basket", 1.0, "robotwin"), + ("Clean/place_bread_skillet", 1.0, "robotwin"), + ("Clean/place_burger_fries", 1.0, "robotwin"), + ("Clean/place_can_basket", 1.0, "robotwin"), + ("Clean/place_cans_plasticbox", 1.0, "robotwin"), + ("Clean/place_container_plate", 1.0, "robotwin"), + ("Clean/place_dual_shoes", 1.0, "robotwin"), + ("Clean/place_empty_cup", 1.0, "robotwin"), + ("Clean/place_fan", 1.0, "robotwin"), + ("Clean/place_mouse_pad", 1.0, "robotwin"), + ("Clean/place_object_basket", 1.0, "robotwin"), + ("Clean/place_object_scale", 1.0, "robotwin"), + ("Clean/place_object_stand", 1.0, "robotwin"), + ("Clean/place_phone_stand", 1.0, "robotwin"), + ("Clean/place_shoe", 1.0, "robotwin"), + ("Clean/press_stapler", 1.0, "robotwin"), + ("Clean/put_bottles_dustbin", 1.0, "robotwin"), + ("Clean/put_object_cabinet", 1.0, "robotwin"), + ("Clean/rotate_qrcode", 1.0, "robotwin"), + ("Clean/scan_object", 1.0, "robotwin"), + ("Clean/shake_bottle", 1.0, "robotwin"), + ("Clean/shake_bottle_horizontally", 1.0, "robotwin"), + ("Clean/stack_blocks_three", 1.0, "robotwin"), + ("Clean/stack_blocks_two", 1.0, "robotwin"), + ("Clean/stack_bowls_three", 1.0, "robotwin"), + ("Clean/stack_bowls_two", 1.0, "robotwin"), + ("Clean/stamp_seal", 1.0, "robotwin"), + ("Clean/turn_switch", 1.0, "robotwin"), + ("Randomized/adjust_bottle", 1.0, "robotwin"), + ("Randomized/beat_block_hammer", 1.0, "robotwin"), + ("Randomized/blocks_ranking_rgb", 1.0, "robotwin"), + ("Randomized/blocks_ranking_size", 1.0, "robotwin"), + ("Randomized/click_alarmclock", 1.0, "robotwin"), + ("Randomized/click_bell", 1.0, "robotwin"), + ("Randomized/dump_bin_bigbin", 1.0, "robotwin"), + ("Randomized/grab_roller", 1.0, "robotwin"), + ("Randomized/handover_block", 1.0, "robotwin"), + ("Randomized/handover_mic", 1.0, "robotwin"), + ("Randomized/hanging_mug", 1.0, "robotwin"), + ("Randomized/lift_pot", 1.0, "robotwin"), + ("Randomized/move_can_pot", 1.0, "robotwin"), + ("Randomized/move_pillbottle_pad", 1.0, "robotwin"), + ("Randomized/move_playingcard_away", 1.0, "robotwin"), + ("Randomized/move_stapler_pad", 1.0, "robotwin"), + ("Randomized/open_laptop", 1.0, "robotwin"), + ("Randomized/open_microwave", 1.0, "robotwin"), + ("Randomized/pick_diverse_bottles", 1.0, "robotwin"), + ("Randomized/pick_dual_bottles", 1.0, "robotwin"), + ("Randomized/place_a2b_left", 1.0, "robotwin"), + ("Randomized/place_a2b_right", 1.0, "robotwin"), + ("Randomized/place_bread_basket", 1.0, "robotwin"), + ("Randomized/place_bread_skillet", 1.0, "robotwin"), + ("Randomized/place_burger_fries", 1.0, "robotwin"), + ("Randomized/place_can_basket", 1.0, "robotwin"), + ("Randomized/place_cans_plasticbox", 1.0, "robotwin"), + ("Randomized/place_container_plate", 1.0, "robotwin"), + ("Randomized/place_dual_shoes", 1.0, "robotwin"), + ("Randomized/place_empty_cup", 1.0, "robotwin"), + ("Randomized/place_fan", 1.0, "robotwin"), + ("Randomized/place_mouse_pad", 1.0, "robotwin"), + ("Randomized/place_object_basket", 1.0, "robotwin"), + ("Randomized/place_object_scale", 1.0, "robotwin"), + ("Randomized/place_object_stand", 1.0, "robotwin"), + ("Randomized/place_phone_stand", 1.0, "robotwin"), + ("Randomized/place_shoe", 1.0, "robotwin"), + ("Randomized/press_stapler", 1.0, "robotwin"), + ("Randomized/put_bottles_dustbin", 1.0, "robotwin"), + ("Randomized/put_object_cabinet", 1.0, "robotwin"), + ("Randomized/rotate_qrcode", 1.0, "robotwin"), + ("Randomized/scan_object", 1.0, "robotwin"), + ("Randomized/shake_bottle", 1.0, "robotwin"), + ("Randomized/shake_bottle_horizontally", 1.0, "robotwin"), + ("Randomized/stack_blocks_three", 1.0, "robotwin"), + ("Randomized/stack_blocks_two", 1.0, "robotwin"), + ("Randomized/stack_bowls_three", 1.0, "robotwin"), + ("Randomized/stack_bowls_two", 1.0, "robotwin"), + ("Randomized/stamp_seal", 1.0, "robotwin"), + ("Randomized/turn_switch", 1.0, "robotwin"), + ], + "robotwin_all_50": [ + ("Clean/adjust_bottle", 1.0, "robotwin50"), + ("Clean/beat_block_hammer", 1.0, "robotwin50"), + ("Clean/blocks_ranking_rgb", 1.0, "robotwin50"), + ("Clean/blocks_ranking_size", 1.0, "robotwin50"), + ("Clean/click_alarmclock", 1.0, "robotwin50"), + ("Clean/click_bell", 1.0, "robotwin50"), + ("Clean/dump_bin_bigbin", 1.0, "robotwin50"), + ("Clean/grab_roller", 1.0, "robotwin50"), + ("Clean/handover_block", 1.0, "robotwin50"), + ("Clean/handover_mic", 1.0, "robotwin50"), + ("Clean/hanging_mug", 1.0, "robotwin50"), + ("Clean/lift_pot", 1.0, "robotwin50"), + ("Clean/move_can_pot", 1.0, "robotwin50"), + ("Clean/move_pillbottle_pad", 1.0, "robotwin50"), + ("Clean/move_playingcard_away", 1.0, "robotwin50"), + ("Clean/move_stapler_pad", 1.0, "robotwin50"), + ("Clean/open_laptop", 1.0, "robotwin50"), + ("Clean/open_microwave", 1.0, "robotwin50"), + ("Clean/pick_diverse_bottles", 1.0, "robotwin50"), + ("Clean/pick_dual_bottles", 1.0, "robotwin50"), + ("Clean/place_a2b_left", 1.0, "robotwin50"), + ("Clean/place_a2b_right", 1.0, "robotwin50"), + ("Clean/place_bread_basket", 1.0, "robotwin50"), + ("Clean/place_bread_skillet", 1.0, "robotwin50"), + ("Clean/place_burger_fries", 1.0, "robotwin50"), + ("Clean/place_can_basket", 1.0, "robotwin50"), + ("Clean/place_cans_plasticbox", 1.0, "robotwin50"), + ("Clean/place_container_plate", 1.0, "robotwin50"), + ("Clean/place_dual_shoes", 1.0, "robotwin50"), + ("Clean/place_empty_cup", 1.0, "robotwin50"), + ("Clean/place_fan", 1.0, "robotwin50"), + ("Clean/place_mouse_pad", 1.0, "robotwin50"), + ("Clean/place_object_basket", 1.0, "robotwin50"), + ("Clean/place_object_scale", 1.0, "robotwin50"), + ("Clean/place_object_stand", 1.0, "robotwin50"), + ("Clean/place_phone_stand", 1.0, "robotwin50"), + ("Clean/place_shoe", 1.0, "robotwin50"), + ("Clean/press_stapler", 1.0, "robotwin50"), + ("Clean/put_bottles_dustbin", 1.0, "robotwin50"), + ("Clean/put_object_cabinet", 1.0, "robotwin50"), + ("Clean/rotate_qrcode", 1.0, "robotwin50"), + ("Clean/scan_object", 1.0, "robotwin50"), + ("Clean/shake_bottle", 1.0, "robotwin50"), + ("Clean/shake_bottle_horizontally", 1.0, "robotwin50"), + ("Clean/stack_blocks_three", 1.0, "robotwin50"), + ("Clean/stack_blocks_two", 1.0, "robotwin50"), + ("Clean/stack_bowls_three", 1.0, "robotwin50"), + ("Clean/stack_bowls_two", 1.0, "robotwin50"), + ("Clean/stamp_seal", 1.0, "robotwin50"), + ("Clean/turn_switch", 1.0, "robotwin50"), + ("Randomized/adjust_bottle", 1.0, "robotwin50"), + ("Randomized/beat_block_hammer", 1.0, "robotwin50"), + ("Randomized/blocks_ranking_rgb", 1.0, "robotwin50"), + ("Randomized/blocks_ranking_size", 1.0, "robotwin50"), + ("Randomized/click_alarmclock", 1.0, "robotwin50"), + ("Randomized/click_bell", 1.0, "robotwin50"), + ("Randomized/dump_bin_bigbin", 1.0, "robotwin50"), + ("Randomized/grab_roller", 1.0, "robotwin50"), + ("Randomized/handover_block", 1.0, "robotwin50"), + ("Randomized/handover_mic", 1.0, "robotwin50"), + ("Randomized/hanging_mug", 1.0, "robotwin50"), + ("Randomized/lift_pot", 1.0, "robotwin50"), + ("Randomized/move_can_pot", 1.0, "robotwin50"), + ("Randomized/move_pillbottle_pad", 1.0, "robotwin50"), + ("Randomized/move_playingcard_away", 1.0, "robotwin50"), + ("Randomized/move_stapler_pad", 1.0, "robotwin50"), + ("Randomized/open_laptop", 1.0, "robotwin50"), + ("Randomized/open_microwave", 1.0, "robotwin50"), + ("Randomized/pick_diverse_bottles", 1.0, "robotwin50"), + ("Randomized/pick_dual_bottles", 1.0, "robotwin50"), + ("Randomized/place_a2b_left", 1.0, "robotwin50"), + ("Randomized/place_a2b_right", 1.0, "robotwin50"), + ("Randomized/place_bread_basket", 1.0, "robotwin50"), + ("Randomized/place_bread_skillet", 1.0, "robotwin50"), + ("Randomized/place_burger_fries", 1.0, "robotwin50"), + ("Randomized/place_can_basket", 1.0, "robotwin50"), + ("Randomized/place_cans_plasticbox", 1.0, "robotwin50"), + ("Randomized/place_container_plate", 1.0, "robotwin50"), + ("Randomized/place_dual_shoes", 1.0, "robotwin50"), + ("Randomized/place_empty_cup", 1.0, "robotwin50"), + ("Randomized/place_fan", 1.0, "robotwin50"), + ("Randomized/place_mouse_pad", 1.0, "robotwin50"), + ("Randomized/place_object_basket", 1.0, "robotwin50"), + ("Randomized/place_object_scale", 1.0, "robotwin50"), + ("Randomized/place_object_stand", 1.0, "robotwin50"), + ("Randomized/place_phone_stand", 1.0, "robotwin50"), + ("Randomized/place_shoe", 1.0, "robotwin50"), + ("Randomized/press_stapler", 1.0, "robotwin50"), + ("Randomized/put_bottles_dustbin", 1.0, "robotwin50"), + ("Randomized/put_object_cabinet", 1.0, "robotwin50"), + ("Randomized/rotate_qrcode", 1.0, "robotwin50"), + ("Randomized/scan_object", 1.0, "robotwin50"), + ("Randomized/shake_bottle", 1.0, "robotwin50"), + ("Randomized/shake_bottle_horizontally", 1.0, "robotwin50"), + ("Randomized/stack_blocks_three", 1.0, "robotwin50"), + ("Randomized/stack_blocks_two", 1.0, "robotwin50"), + ("Randomized/stack_bowls_three", 1.0, "robotwin50"), + ("Randomized/stack_bowls_two", 1.0, "robotwin50"), + ("Randomized/stamp_seal", 1.0, "robotwin50"), + ("Randomized/turn_switch", 1.0, "robotwin50"), + ], + "robotwin": [ + ("adjust_bottle", 1.0, "robotwin"), + ("beat_block_hammer", 1.0, "robotwin"), + ("blocks_ranking_rgb", 1.0, "robotwin"), + ("blocks_ranking_size", 1.0, "robotwin"), + ("click_alarmclock", 1.0, "robotwin"), + ("click_bell", 1.0, "robotwin"), + ("dump_bin_bigbin", 1.0, "robotwin"), + ("grab_roller", 1.0, "robotwin"), + ("handover_block", 1.0, "robotwin"), + ("handover_mic", 1.0, "robotwin"), + ("hanging_mug", 1.0, "robotwin"), + ("lift_pot", 1.0, "robotwin"), + ("move_can_pot", 1.0, "robotwin"), + ("move_pillbottle_pad", 1.0, "robotwin"), + ("move_playingcard_away", 1.0, "robotwin"), + ("move_stapler_pad", 1.0, "robotwin"), + ("open_laptop", 1.0, "robotwin"), + ("open_microwave", 1.0, "robotwin"), + ("pick_diverse_bottles", 1.0, "robotwin"), + ("pick_dual_bottles", 1.0, "robotwin"), + ("place_a2b_left", 1.0, "robotwin"), + ("place_a2b_right", 1.0, "robotwin"), + ("place_bread_basket", 1.0, "robotwin"), + ("place_bread_skillet", 1.0, "robotwin"), + ("place_burger_fries", 1.0, "robotwin"), + ("place_can_basket", 1.0, "robotwin"), + ("place_cans_plasticbox", 1.0, "robotwin"), + ("place_container_plate", 1.0, "robotwin"), + ("place_dual_shoes", 1.0, "robotwin"), + ("place_empty_cup", 1.0, "robotwin"), + ("place_fan", 1.0, "robotwin"), + ("place_mouse_pad", 1.0, "robotwin"), + ("place_object_basket", 1.0, "robotwin"), + ("place_object_scale", 1.0, "robotwin"), + ("place_object_stand", 1.0, "robotwin"), + ("place_phone_stand", 1.0, "robotwin"), + ("place_shoe", 1.0, "robotwin"), + ("press_stapler", 1.0, "robotwin"), + ("put_bottles_dustbin", 1.0, "robotwin"), + ("put_object_cabinet", 1.0, "robotwin"), + ("rotate_qrcode", 1.0, "robotwin"), + ("scan_object", 1.0, "robotwin"), + ("shake_bottle", 1.0, "robotwin"), + ("shake_bottle_horizontally", 1.0, "robotwin"), + ("stack_blocks_three", 1.0, "robotwin"), + ("stack_blocks_two", 1.0, "robotwin"), + ("stack_bowls_three", 1.0, "robotwin"), + ("stack_bowls_two", 1.0, "robotwin"), + ("stamp_seal", 1.0, "robotwin"), + ("turn_switch", 1.0, "robotwin"), + ], + + "robotwin_task1": [ + ("adjust_bottle", 1.0, "robotwin"), + ], + "robotwin_task2": [ + ("place_a2b_left", 1.0, "robotwin"), + ("place_a2b_right", 1.0, "robotwin"), + ], + + "multi_robot": [ + ("LEROBOT_LIBERO_DATA/libero_10_no_noops_1.0.0_lerobot", 1.0, "libero_franka"), + # ("OXE_LEROBOT_DATASET/bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"), + ], + + "calvin_task_D_D_v3.0": [ + ("calvin_task_D_D_v3.0", 1.0, "libero_franka"), + ], + + "vla_arena_L0_S": [ + ("VLA_Arena_L0_S_lerobot_openpi", 1.0, "vla_arena_franka"), + ], + "vla_arena_L0_M": [ + ("VLA_Arena_L0_M_lerobot_openpi", 1.0, "vla_arena_franka"), + ], + "vla_arena_L0_L": [ + ("VLA_Arena_L0_L_lerobot_openpi", 1.0, "vla_arena_franka"), + ], +} diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/registry.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..581b926699a9ca9240d361ddb972db00f4867bf6 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/registry.py @@ -0,0 +1,149 @@ +""" +Centralized registry that auto-discovers benchmark-specific data configs from +``examples/*/train_files/data_registry/`` and merges them with the base +registries defined in this package. + +Three registries are maintained: + +* ``DATASET_NAMED_MIXTURES`` – mixture_name → [(dataset, weight, robot_type)] +* ``ROBOT_TYPE_CONFIG_MAP`` – robot_type → DataConfig instance +* ``ROBOT_TYPE_TO_EMBODIMENT_TAG`` – robot_type → EmbodimentTag + +``ROBOT_TYPE_TO_EMBODIMENT_TAG`` is **derived** from ``ROBOT_TYPE_CONFIG_MAP`` +by reading each DataConfig class's ``embodiment_tag`` classvar (Proposal A). +Classes without the classvar fall back to ``EmbodimentTag.NEW_EMBODIMENT``. +Legacy bench files exposing their own ``ROBOT_TYPE_TO_EMBODIMENT_TAG`` dict are +still honored as overrides for backward compatibility. + + +Usage:: + + from starVLA.dataloader.gr00t_lerobot.registry import ( + ROBOT_TYPE_CONFIG_MAP, + ROBOT_TYPE_TO_EMBODIMENT_TAG, + DATASET_NAMED_MIXTURES, + ) +""" + +from __future__ import annotations + +import importlib +import importlib.util +import logging +import sys +from pathlib import Path + +# Base registries (kept as fallback / seed values) +from starVLA.dataloader.gr00t_lerobot.data_config import ( + ROBOT_TYPE_CONFIG_MAP as _BASE_CONFIG_MAP, +) +from starVLA.dataloader.gr00t_lerobot.embodiment_tags import ( + EmbodimentTag, # re-export for convenience +) +from starVLA.dataloader.gr00t_lerobot.mixtures import ( + DATASET_NAMED_MIXTURES as _BASE_MIXTURES, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Mutable copies – will be extended by discovered modules +# --------------------------------------------------------------------------- +ROBOT_TYPE_CONFIG_MAP: dict = dict(_BASE_CONFIG_MAP) +DATASET_NAMED_MIXTURES: dict = dict(_BASE_MIXTURES) + +# Legacy explicit overrides (rarely needed; prefer the classvar on DataConfig). +_LEGACY_TAG_OVERRIDES: dict = {} + + +def _derive_tag_map() -> dict: + """Build robot_type -> EmbodimentTag from ROBOT_TYPE_CONFIG_MAP. + + Reads each DataConfig instance's ``embodiment_tag`` classvar. Falls back to + ``EmbodimentTag.NEW_EMBODIMENT`` if the classvar is missing. Legacy + overrides from bench files take precedence. + """ + out: dict = {} + for rt, cfg in ROBOT_TYPE_CONFIG_MAP.items(): + tag = getattr(cfg, "embodiment_tag", None) + if tag is None: + logger.warning( + "[registry] DataConfig for robot_type=%r has no `embodiment_tag` " + "classvar; defaulting to EmbodimentTag.NEW_EMBODIMENT.", rt + ) + tag = EmbodimentTag.NEW_EMBODIMENT + out[rt] = tag + out.update(_LEGACY_TAG_OVERRIDES) + return out + + +ROBOT_TYPE_TO_EMBODIMENT_TAG: dict = {} # populated by discover_and_merge() + +# --------------------------------------------------------------------------- +# Discovery logic +# --------------------------------------------------------------------------- +_REGISTRY_DIR_NAME = "data_registry" +_DISCOVERED = False + + +def _find_registry_dirs() -> list[Path]: + """Return all ``examples/*/train_files/data_registry/`` directories.""" + # Walk up from this file to the repo root + # registry.py is at starVLA/dataloader/gr00t_lerobot/registry.py + # parents: [0]=gr00t_lerobot, [1]=dataloader, [2]=starVLA(pkg), [3]=repo root + repo_root = Path(__file__).resolve().parents[3] + examples_dir = repo_root / "examples" + if not examples_dir.is_dir(): + return [] + dirs: list[Path] = [] + for bench_dir in sorted(examples_dir.iterdir()): + registry_dir = bench_dir / "train_files" / _REGISTRY_DIR_NAME + if registry_dir.is_dir(): + dirs.append(registry_dir) + return dirs + + +def _load_module_from_path(module_name: str, file_path: Path): + """Import a Python file as a module with the given name.""" + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None or spec.loader is None: + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[module_name] = mod + spec.loader.exec_module(mod) + return mod + + +def discover_and_merge() -> None: + """Scan ``examples/*/train_files/data_registry/`` and merge into global registries.""" + global _DISCOVERED, ROBOT_TYPE_TO_EMBODIMENT_TAG + if _DISCOVERED: + return + _DISCOVERED = True + + for registry_dir in _find_registry_dirs(): + bench_name = registry_dir.parents[1].name # examples//train_files/data_registry + prefix = f"_data_registry_{bench_name}" + + # --- data_config.py (may contain all three registries) --- + cfg_file = registry_dir / "data_config.py" + if cfg_file.is_file(): + mod = _load_module_from_path(f"{prefix}.data_config", cfg_file) + if mod: + if hasattr(mod, "ROBOT_TYPE_CONFIG_MAP"): + ROBOT_TYPE_CONFIG_MAP.update(mod.ROBOT_TYPE_CONFIG_MAP) + logger.debug(f"[registry] Loaded data_config from {bench_name}: {list(mod.ROBOT_TYPE_CONFIG_MAP.keys())}") + # Legacy: bench file may still expose an explicit dict override. + if hasattr(mod, "ROBOT_TYPE_TO_EMBODIMENT_TAG"): + _LEGACY_TAG_OVERRIDES.update(mod.ROBOT_TYPE_TO_EMBODIMENT_TAG) + logger.debug(f"[registry] Legacy embodiment_tag overrides from {bench_name}: {list(mod.ROBOT_TYPE_TO_EMBODIMENT_TAG.keys())}") + if hasattr(mod, "DATASET_NAMED_MIXTURES"): + DATASET_NAMED_MIXTURES.update(mod.DATASET_NAMED_MIXTURES) + logger.debug(f"[registry] Loaded mixtures from {bench_name} (data_config): {list(mod.DATASET_NAMED_MIXTURES.keys())}") + + # Re-derive after all benches are loaded so classvars + legacy overrides combine. + ROBOT_TYPE_TO_EMBODIMENT_TAG = _derive_tag_map() + + +# Run discovery on first import +discover_and_merge() diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/schema.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..64519e56a3c59e5d08c8f8f6370f640061859b4d --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/schema.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 enum import Enum +from typing import Optional + +from numpydantic import NDArray +from pydantic import BaseModel, Field, field_serializer + +from .embodiment_tags import EmbodimentTag + +# Common schema + + +class RotationType(Enum): + """Type of rotation representation""" + + AXIS_ANGLE = "axis_angle" + QUATERNION = "quaternion" + ROTATION_6D = "rotation_6d" + MATRIX = "matrix" + EULER_ANGLES_RPY = "euler_angles_rpy" + EULER_ANGLES_RYP = "euler_angles_ryp" + EULER_ANGLES_PRY = "euler_angles_pry" + EULER_ANGLES_PYR = "euler_angles_pyr" + EULER_ANGLES_YRP = "euler_angles_yrp" + EULER_ANGLES_YPR = "euler_angles_ypr" + + +# LeRobot schema + + +class LeRobotModalityField(BaseModel): + """Metadata for a LeRobot modality field.""" + + original_key: Optional[str] = Field( + default=None, + description="The original key of the modality in the LeRobot dataset", + ) + + +class LeRobotStateActionMetadata(LeRobotModalityField): + """Metadata for a LeRobot modality.""" + + start: int = Field( + ..., + description="The start index of the modality in the concatenated state/action vector", + ) + end: int = Field( + ..., + description="The end index of the modality in the concatenated state/action vector", + ) + rotation_type: Optional[RotationType] = Field( + default=None, description="The type of rotation for the modality" + ) + absolute: bool = Field(default=True, description="Whether the modality is absolute") + dtype: str = Field( + default="float64", + description="The data type of the modality. Defaults to float64.", + ) + range: Optional[tuple[float, float]] = Field( + default=None, + description="The range of the modality, if applicable. Defaults to None.", + ) + original_key: Optional[str] = Field( + default=None, + description="The original key of the modality in the LeRobot dataset.", + ) + + +class LeRobotStateMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot state modality.""" + + original_key: Optional[str] = Field( + default="observation.state", # LeRobot convention for states + description="The original key of the state modality in the LeRobot dataset", + ) + + +class LeRobotActionMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot action modality.""" + + original_key: Optional[str] = Field( + default="action", # LeRobot convention for actions + description="The original key of the action modality in the LeRobot dataset", + ) + + +class LeRobotModalityMetadata(BaseModel): + """Metadata for a LeRobot modality.""" + + state: dict[str, LeRobotStateMetadata] = Field( + ..., + description="The metadata for the state modality. The keys are the names of each split of the state vector.", + ) + action: dict[str, LeRobotActionMetadata] = Field( + ..., + description="The metadata for the action modality. The keys are the names of each split of the action vector.", + ) + video: dict[str, LeRobotModalityField] = Field( + ..., + description="The metadata for the video modality. The keys are the new names of each video modality.", + ) + annotation: Optional[dict[str, LeRobotModalityField]] = Field( + default=None, + description="The metadata for the annotation modality. The keys are the new names of each annotation modality.", + ) + + def get_key_meta(self, key: str) -> LeRobotModalityField: + """Get the metadata for a key in the LeRobot modality metadata. + + Args: + key (str): The key to get the metadata for. + + Returns: + LeRobotModalityField: The metadata for the key. + + Example: + lerobot_modality_meta = LeRobotModalityMetadata.model_validate(U.load_json(modality_meta_path)) + lerobot_modality_meta.get_key_meta("state.joint_shoulder_y") + lerobot_modality_meta.get_key_meta("video.main_camera") + lerobot_modality_meta.get_key_meta("annotation.human.action.task_description") + """ + split_key = key.split(".") + modality = split_key[0] + subkey = ".".join(split_key[1:]) + if modality == "state": + if subkey not in self.state: + raise ValueError( + f"Key: {key}, state key {subkey} not found in metadata, available state keys: {self.state.keys()}" + ) + return self.state[subkey] + elif modality == "action": + if subkey not in self.action: + raise ValueError( + f"Key: {key}, action key {subkey} not found in metadata, available action keys: {self.action.keys()}" + ) + return self.action[subkey] + elif modality == "video": + if subkey not in self.video: + raise ValueError( + f"Key: {key}, video key {subkey} not found in metadata, available video keys: {self.video.keys()}" + ) + return self.video[subkey] + elif modality == "annotation": + assert ( + self.annotation is not None + ), "Trying to get annotation metadata for a dataset with no annotations" + if subkey not in self.annotation: + raise ValueError( + f"Key: {key}, annotation key {subkey} not found in metadata, available annotation keys: {self.annotation.keys()}" + ) + return self.annotation[subkey] + else: + raise ValueError(f"Key: {key}, unexpected modality: {modality}") + + +# Dataset schema (parsed from LeRobot schema and simplified) + + +class DatasetStatisticalValues(BaseModel): + max: NDArray = Field(..., description="Maximum values") + min: NDArray = Field(..., description="Minimum values") + mean: NDArray = Field(..., description="Mean values") + std: NDArray = Field(..., description="Standard deviation") + q01: NDArray = Field(..., description="1st percentile values") + q99: NDArray = Field(..., description="99th percentile values") + + @field_serializer("*", when_used="json") + def serialize_ndarray(self, v: NDArray) -> list[float]: + return v.tolist() # type: ignore + + +class DatasetStatistics(BaseModel): + state: dict[str, DatasetStatisticalValues] = Field(..., description="Statistics of the state") + action: dict[str, DatasetStatisticalValues] = Field(..., description="Statistics of the action") + + +class VideoMetadata(BaseModel): + """Metadata of the video modality""" + + resolution: tuple[int, int] = Field(..., description="Resolution of the video") + channels: int = Field(..., description="Number of channels in the video", gt=0) + fps: float = Field(..., description="Frames per second", gt=0) + + +class StateActionMetadata(BaseModel): + absolute: bool = Field(..., description="Whether the state or action is absolute") + rotation_type: Optional[RotationType] = Field(None, description="Type of rotation, if any") + shape: tuple[int, ...] = Field(..., description="Shape of the state or action") + continuous: bool = Field(..., description="Whether the state or action is continuous") + + +class DatasetModalities(BaseModel): + video: dict[str, VideoMetadata] = Field(..., description="Metadata of the video") + state: dict[str, StateActionMetadata] = Field(..., description="Metadata of the state") + action: dict[str, StateActionMetadata] = Field(..., description="Metadata of the action") + + +class DatasetMetadata(BaseModel): + """Metadata of the trainable dataset + + Changes: + - Update to use the new RawCommitHashMetadataMetadata_V1_2 + """ + + statistics: DatasetStatistics = Field(..., description="Statistics of the dataset") + modalities: DatasetModalities = Field(..., description="Metadata of the modalities") + embodiment_tag: EmbodimentTag = Field(..., description="Embodiment tag of the dataset") \ No newline at end of file diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/__init__.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bf22b5e4ca3d1c7937a25234cbf08e2644593587 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/__init__.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 .base import ( + ComposedModalityTransform, + InvertibleModalityTransform, + ModalityTransform, +) +from .concat import ConcatTransform +# from .state_action import ( +# StateActionDropout, +# StateActionPerturbation, +# StateActionSinCosTransform, +# StateActionToTensor, +# StateActionTransform, +# ) +from .video import ( + VideoColorJitter, + VideoCrop, + VideoGrayscale, + VideoHorizontalFlip, + VideoRandomGrayscale, + VideoRandomPosterize, + VideoRandomRotation, + VideoResize, + VideoToNumpy, + VideoToTensor, + VideoTransform, +) diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/base.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/base.py new file mode 100644 index 0000000000000000000000000000000000000000..aac88559af98fa23f34fbb9135775d0819c281ef --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/base.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 abc import ABC, abstractmethod +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr + +from ..schema import DatasetMetadata + + +class ModalityTransform(BaseModel, ABC): + """ + Abstract class for transforming data modalities, e.g. video frame augmentation or action normalization. + """ + + apply_to: list[str] = Field(..., description="The keys to apply the transform to.") + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + _dataset_metadata: DatasetMetadata | None = PrivateAttr(default=None) + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @property + def dataset_metadata(self) -> DatasetMetadata: + assert ( + self._dataset_metadata is not None + ), "Dataset metadata is not set. Please call set_metadata() before calling apply()." + return self._dataset_metadata + + @dataset_metadata.setter + def dataset_metadata(self, value: DatasetMetadata): + self._dataset_metadata = value + + def set_metadata(self, dataset_metadata: DatasetMetadata): + """ + Set the dataset metadata. This is useful for transforms that need to know the dataset metadata, e.g. to normalize actions. + Subclasses can override this method if they need to do something more complex. + """ + self.dataset_metadata = dataset_metadata + + def __call__(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply the transformation to the data corresponding to target_keys and return the processed data. + + Args: + data (dict[str, Any]): The data to transform. + example: data = { + "video.image_side_0": np.ndarray, + "action.eef_position": np.ndarray, + ... + } + + Returns: + dict[str, Any]: The transformed data. + example: transformed_data = { + "video.image_side_0": np.ndarray, + "action.eef_position": torch.Tensor, # Normalized and converted to tensor + ... + } + """ + return self.apply(data) + + @abstractmethod + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply the transformation to the data corresponding to keys matching the `apply_to` regular expression and return the processed data.""" + + def train(self): + self.training = True + + def eval(self): + self.training = False + + +class InvertibleModalityTransform(ModalityTransform): + @abstractmethod + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Reverse the transformation to the data corresponding to keys matching the `apply_to` regular expression and return the processed data.""" + + +class ComposedModalityTransform(ModalityTransform): + """Compose multiple modality transforms.""" + + transforms: list[ModalityTransform] = Field(..., description="The transforms to compose.") + apply_to: list[str] = Field( + default_factory=list, description="Will be ignored for composed transforms." + ) + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + + model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True) + + def set_metadata(self, dataset_metadata: DatasetMetadata): + for transform in self.transforms: + transform.set_metadata(dataset_metadata) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for i, transform in enumerate(self.transforms): + try: + data = transform(data) + except Exception as e: + raise ValueError(f"Error applying transform {i} to data: {e}") from e + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for i, transform in enumerate(reversed(self.transforms)): + if isinstance(transform, InvertibleModalityTransform): + try: + data = transform.unapply(data) + except Exception as e: + step = len(self.transforms) - i - 1 + raise ValueError(f"Error unapplying transform {step} to data: {e}") from e + return data + + def train(self): + for transform in self.transforms: + transform.train() + + def eval(self): + for transform in self.transforms: + transform.eval() diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/concat.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/concat.py new file mode 100644 index 0000000000000000000000000000000000000000..cf8eea4c77fc163ecdb0d25aeca26a2cde99f8c4 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/concat.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 Optional + +import numpy as np +import torch +from pydantic import Field + +from ..schema import DatasetMetadata, StateActionMetadata +from .base import InvertibleModalityTransform + + +class ConcatTransform(InvertibleModalityTransform): + """ + Concatenate the keys according to specified order. + """ + + # -- We inherit from ModalityTransform, so we keep apply_to as well -- + apply_to: list[str] = Field( + default_factory=list, description="Not used in this transform, kept for compatibility." + ) + + video_concat_order: list[str] = Field( + ..., + description="Concatenation order for each video modality. " + "Format: ['video.ego_view_pad_res224_freq20', ...]", + ) + + state_concat_order: Optional[list[str]] = Field( + default=None, + description="Concatenation order for each state modality. " + "Format: ['state.position', 'state.velocity', ...].", + ) + + action_concat_order: Optional[list[str]] = Field( + default=None, + description="Concatenation order for each action modality. " + "Format: ['action.position', 'action.velocity', ...].", + ) + + action_dims: dict[str, int] = Field( + default_factory=dict, + description="The dimensions of the action keys.", + ) + state_dims: dict[str, int] = Field( + default_factory=dict, + description="The dimensions of the state keys.", + ) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = { + "apply_to", + "video_concat_order", + "state_concat_order", + "action_concat_order", + } + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + def apply(self, data: dict) -> dict: + grouped_keys = {} + for key in data.keys(): + try: + modality, _ = key.split(".") + except: # noqa: E722 + ### Handle language annotation special case + if "annotation" in key: + modality = "language" + else: + modality = "others" + if modality not in grouped_keys: + grouped_keys[modality] = [] + grouped_keys[modality].append(key) + + if "video" in grouped_keys: + # Check if keys in video_concat_order, state_concat_order, action_concat_order are + # ineed contained in the data. If not, then the keys are misspecified + video_keys = grouped_keys["video"] + assert self.video_concat_order is not None, f"{self.video_concat_order=}, {video_keys=}" + assert all( + item in video_keys for item in self.video_concat_order + ), f"keys in video_concat_order are misspecified, \n{video_keys=}, \n{self.video_concat_order=}" + + # Process each video view + unsqueezed_videos = [] + for video_key in self.video_concat_order: + video_data = data.pop(video_key) + unsqueezed_video = np.expand_dims( + video_data, axis=-4 + ) # [..., H, W, C] -> [..., 1, H, W, C] + unsqueezed_videos.append(unsqueezed_video) + # Concatenate along the new axis + unsqueezed_video = np.concatenate(unsqueezed_videos, axis=-4) # [..., V, H, W, C] + + # Video + data["video"] = unsqueezed_video + + # "state" + if "state" in grouped_keys: + state_keys = grouped_keys["state"] + assert self.state_concat_order is not None, f"{self.state_concat_order=}" + assert all( + item in state_keys for item in self.state_concat_order + ), f"keys in state_concat_order are misspecified, \n{state_keys=}, \n{self.state_concat_order=}" + # Check the state dims + for key in self.state_concat_order: + target_shapes = [self.state_dims[key]] + if self.is_rotation_key(key): + target_shapes.append(6) # Allow for rotation_6d + # if key in ["state.right_arm", "state.right_hand"]: + target_shapes.append(self.state_dims[key] * 2) # Allow for sin-cos transform + assert ( + data[key].shape[-1] in target_shapes + ), f"State dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}" + # Concatenate the state keys + # We'll have StateActionToTensor before this transform, so here we use torch.cat + data["state"] = torch.cat( + [data.pop(key) for key in self.state_concat_order], dim=-1 + ) # [T, D_state] + + if "action" in grouped_keys: + action_keys = grouped_keys["action"] + assert self.action_concat_order is not None, f"{self.action_concat_order=}" + # Check if all keys in concat_order are present + assert set(self.action_concat_order) == set( + action_keys + ), f"{set(self.action_concat_order)=}, {set(action_keys)=}" + # Record the action dims + for key in self.action_concat_order: + target_shapes = [self.action_dims[key]] + if self.is_rotation_key(key): + target_shapes.append(3) # Allow for axis angle + assert ( + self.action_dims[key] == data[key].shape[-1] + ), f"Action dim mismatch for {key=}, {self.action_dims[key]=}, {data[key].shape[-1]=}" + # Concatenate the action keys + # We'll have StateActionToTensor before this transform, so here we use torch.cat + data["action"] = torch.cat( + [data.pop(key) for key in self.action_concat_order], dim=-1 + ) # [T, D_action] + + return data + + def unapply(self, data: dict) -> dict: + start_dim = 0 + assert "action" in data, f"{data.keys()=}" + # For those dataset without actions (LAPA), we'll never run unapply + assert self.action_concat_order is not None, f"{self.action_concat_order=}" + action_tensor = data.pop("action") + for key in self.action_concat_order: + if key not in self.action_dims: + raise ValueError(f"Action dim {key} not found in action_dims.") + end_dim = start_dim + self.action_dims[key] + data[key] = action_tensor[..., start_dim:end_dim] + start_dim = end_dim + if "state" in data: + assert self.state_concat_order is not None, f"{self.state_concat_order=}" + start_dim = 0 + state_tensor = data.pop("state") + for key in self.state_concat_order: + end_dim = start_dim + self.state_dims[key] + data[key] = state_tensor[..., start_dim:end_dim] + start_dim = end_dim + return data + + def __call__(self, data: dict) -> dict: + return self.apply(data) + + def get_modality_metadata(self, key: str) -> StateActionMetadata: + modality, subkey = key.split(".") + assert self.dataset_metadata is not None, "Metadata not set" + modality_config = getattr(self.dataset_metadata.modalities, modality) + assert subkey in modality_config, f"{subkey=} not found in {modality_config=}" + assert isinstance( + modality_config[subkey], StateActionMetadata + ), f"Expected {StateActionMetadata} for {subkey=}, got {type(modality_config[subkey])=}" + return modality_config[subkey] + + def get_state_action_dims(self, key: str) -> int: + """Get the dimension of a state or action key from the dataset metadata.""" + modality_config = self.get_modality_metadata(key) + shape = modality_config.shape + assert len(shape) == 1, f"{shape=}" + return shape[0] + + def is_rotation_key(self, key: str) -> bool: + modality_config = self.get_modality_metadata(key) + return modality_config.rotation_type is not None + + def set_metadata(self, dataset_metadata: DatasetMetadata): + """Set the metadata and compute the dimensions of the state and action keys.""" + super().set_metadata(dataset_metadata) + # Pre-compute the dimensions of the state and action keys + if self.action_concat_order is not None: + for key in self.action_concat_order: + self.action_dims[key] = self.get_state_action_dims(key) + if self.state_concat_order is not None: + for key in self.state_concat_order: + self.state_dims[key] = self.get_state_action_dims(key) diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/state_action.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/state_action.py new file mode 100644 index 0000000000000000000000000000000000000000..51190a87e5c5092953802999baf986488f875311 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/state_action.py @@ -0,0 +1,627 @@ +# Modified in the ReTime research artifact; see NOTICE. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import functools +import random +from typing import Any, ClassVar + +import numpy as np +import torch +from pydantic import Field, PrivateAttr, field_validator, model_validator + +try: + import pytorch3d.transforms as pt + _PYTORCH3D_IMPORT_ERROR = None +except ModuleNotFoundError as exc: # PyTorch3D has no official Windows wheels. + if exc.name is None or not exc.name.startswith("pytorch3d"): + raise + pt = None + _PYTORCH3D_IMPORT_ERROR = exc + +from ..schema import DatasetMetadata, RotationType, StateActionMetadata +from .base import InvertibleModalityTransform, ModalityTransform + + +class RotationTransform: + """Adapted from https://github.com/real-stanford/diffusion_policy/blob/548a52bbb105518058e27bf34dcf90bf6f73681a/diffusion_policy/model/common/rotation_transformer.py""" + + valid_reps = ["axis_angle", "euler_angles", "quaternion", "rotation_6d", "matrix"] + + def __init__(self, from_rep="axis_angle", to_rep="rotation_6d"): + """ + Valid representations + + Always use matrix as intermediate representation. + """ + if pt is None: + raise ImportError( + "Rotation conversion requires pytorch3d. Install pytorch3d or use a " + "data configuration without target_rotations (the LIBERO configuration " + "does not require it)." + ) from _PYTORCH3D_IMPORT_ERROR + + if from_rep.startswith("euler_angles"): + from_convention = from_rep.split("_")[-1] + from_rep = "euler_angles" + from_convention = from_convention.replace("r", "X").replace("p", "Y").replace("y", "Z") + else: + from_convention = None + if to_rep.startswith("euler_angles"): + to_convention = to_rep.split("_")[-1] + to_rep = "euler_angles" + to_convention = to_convention.replace("r", "X").replace("p", "Y").replace("y", "Z") + else: + to_convention = None + assert from_rep != to_rep, f"from_rep and to_rep cannot be the same: {from_rep}" + assert from_rep in self.valid_reps, f"Invalid from_rep: {from_rep}" + assert to_rep in self.valid_reps, f"Invalid to_rep: {to_rep}" + + forward_funcs = list() + inverse_funcs = list() + + if from_rep != "matrix": + funcs = [getattr(pt, f"{from_rep}_to_matrix"), getattr(pt, f"matrix_to_{from_rep}")] + if from_convention is not None: + funcs = [functools.partial(func, convention=from_convention) for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + if to_rep != "matrix": + funcs = [getattr(pt, f"matrix_to_{to_rep}"), getattr(pt, f"{to_rep}_to_matrix")] + if to_convention is not None: + funcs = [functools.partial(func, convention=to_convention) for func in funcs] + forward_funcs.append(funcs[0]) + inverse_funcs.append(funcs[1]) + + inverse_funcs = inverse_funcs[::-1] + + self.forward_funcs = forward_funcs + self.inverse_funcs = inverse_funcs + + @staticmethod + def _apply_funcs(x: torch.Tensor, funcs: list) -> torch.Tensor: + assert isinstance(x, torch.Tensor) + for func in funcs: + x = func(x) + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + return self._apply_funcs(x, self.forward_funcs) + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + return self._apply_funcs(x, self.inverse_funcs) + + +class Normalizer: + valid_modes = ["q99", "mean_std", "min_max", "binary"] + + def __init__(self, mode: str, statistics: dict, binary_threshold: float = 0.5): + self.mode = mode + self.statistics = statistics + self.binary_threshold = binary_threshold + for key, value in self.statistics.items(): + self.statistics[key] = torch.tensor(value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Normalize the tensor + if self.mode == "q99": + # Range of q99 is [-1, 1] + q01 = self.statistics["q01"].to(x.dtype) + q99 = self.statistics["q99"].to(x.dtype) + + # In the case of q01 == q99, the normalization will be undefined + # So we set the normalized values to the original values + mask = q01 != q99 + normalized = torch.zeros_like(x) + + # Normalize the values where q01 != q99 + # Formula: 2 * (x - q01) / (q99 - q01) - 1 + normalized[..., mask] = (x[..., mask] - q01[..., mask]) / ( + q99[..., mask] - q01[..., mask] + ) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + # Set the normalized values to the original values where q01 == q99 + normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + + # Clip the normalized values to be between -1 and 1 + normalized = torch.clamp(normalized, -2.2, 2.2) + + elif self.mode == "mean_std": + # Range of mean_std is not fixed, but can be positive or negative + mean = self.statistics["mean"].to(x.dtype) + std = self.statistics["std"].to(x.dtype) + + # In the case of std == 0, the normalization will be undefined + # So we set the normalized values to the original values + mask = std != 0 + normalized = torch.zeros_like(x) + + # Normalize the values where std != 0 + # Formula: (x - mean) / std + normalized[..., mask] = (x[..., mask] - mean[..., mask]) / std[..., mask] + + # Set the normalized values to the original values where std == 0 + normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + + elif self.mode == "min_max": + # Range of min_max is [-1, 1] + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + + # In the case of min == max, the normalization will be undefined + # So we set the normalized values to the original values + mask = min != max + normalized = torch.zeros_like(x) + + # Normalize the values where min != max + # Formula: 2 * (x - min) / (max - min) - 1 + normalized[..., mask] = (x[..., mask] - min[..., mask]) / ( + max[..., mask] - min[..., mask] + ) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + # Set the normalized values to the original values where min == max + # normalized[..., ~mask] = x[..., ~mask].to(x.dtype) + # Set the normalized values to 0 where min == max + normalized[..., ~mask] = 0 + + elif self.mode == "scale": + # Range of scale is [0, 1] + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + abs_max = torch.max(torch.abs(min), torch.abs(max)) + mask = abs_max != 0 + normalized = torch.zeros_like(x) + normalized[..., mask] = x[..., mask] / abs_max[..., mask] + normalized[..., ~mask] = 0 + + elif self.mode == "binary": + # Range of binary is [0, 1] + normalized = (x > self.binary_threshold).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + return normalized + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + if self.mode == "q99": + q01 = self.statistics["q01"].to(x.dtype) + q99 = self.statistics["q99"].to(x.dtype) + return (x + 1) / 2 * (q99 - q01) + q01 + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(x.dtype) + std = self.statistics["std"].to(x.dtype) + return x * std + mean + elif self.mode == "min_max": + min = self.statistics["min"].to(x.dtype) + max = self.statistics["max"].to(x.dtype) + return (x + 1) / 2 * (max - min) + min + elif self.mode == "binary": + return (x >= self.binary_threshold).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + +class StateActionToTensor(InvertibleModalityTransform): + """ + Transforms states and actions to tensors. + """ + + input_dtypes: dict[str, np.dtype] = Field( + default_factory=dict, description="The input dtypes for each state key." + ) + output_dtypes: dict[str, torch.dtype] = Field( + default_factory=dict, description="The output dtypes for each state key." + ) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to"} + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + @field_validator("input_dtypes", "output_dtypes", mode="before") + def validate_dtypes(cls, v): + for key, dtype in v.items(): + if isinstance(dtype, str): + if dtype.startswith("torch."): + dtype_split = dtype.split(".")[-1] + v[key] = getattr(torch, dtype_split) + elif dtype.startswith("np.") or dtype.startswith("numpy."): + dtype_split = dtype.split(".")[-1] + v[key] = np.dtype(dtype_split) + else: + raise ValueError(f"Invalid dtype: {dtype}") + return v + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + value = data[key] + assert isinstance( + value, np.ndarray + ), f"Unexpected input type: {type(value)}. Expected type: {np.ndarray}" + data[key] = torch.from_numpy(value) + if key in self.output_dtypes: + data[key] = data[key].to(self.output_dtypes[key]) + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + value = data[key] + assert isinstance( + value, torch.Tensor + ), f"Unexpected input type: {type(value)}. Expected type: {torch.Tensor}" + data[key] = value.numpy() + if key in self.input_dtypes: + data[key] = data[key].astype(self.input_dtypes[key]) + return data + + +class StateActionTransform(InvertibleModalityTransform): + """ + Class for state or action transform. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + normalization_modes (dict[str, str]): The normalization modes for each state key. + If a state key in apply_to is not present in the dictionary, it will not be normalized. + target_rotations (dict[str, str]): The target representations for each state key. + If a state key in apply_to is not present in the dictionary, it will not be rotated. + """ + + # Configurable attributes + apply_to: list[str] = Field(..., description="The keys in the modality to load and transform.") + normalization_modes: dict[str, str] = Field( + default_factory=dict, description="The normalization modes for each state key." + ) + target_rotations: dict[str, str] = Field( + default_factory=dict, description="The target representations for each state key." + ) + normalization_statistics: dict[str, dict] = Field( + default_factory=dict, description="The statistics for each state key." + ) + binary_threshold: float = Field( + default=0.5, description="Threshold for binary normalization mode." + ) + modality_metadata: dict[str, StateActionMetadata] = Field( + default_factory=dict, description="The modality metadata for each state key." + ) + + # Model variables + _rotation_transformers: dict[str, RotationTransform] = PrivateAttr(default_factory=dict) + _normalizers: dict[str, Normalizer] = PrivateAttr(default_factory=dict) + _input_dtypes: dict[str, np.dtype | torch.dtype] = PrivateAttr(default_factory=dict) + + # Model constants + _DEFAULT_MIN_MAX_STATISTICS: ClassVar[dict] = { + "rotation_6d": { + "min": [-1, -1, -1, -1, -1, -1], + "max": [1, 1, 1, 1, 1, 1], + }, + "euler_angles": { + "min": [-np.pi, -np.pi, -np.pi], + "max": [np.pi, np.pi, np.pi], + }, + "quaternion": { + "min": [-1, -1, -1, -1], + "max": [1, 1, 1, 1], + }, + "axis_angle": { + "min": [-np.pi, -np.pi, -np.pi], + "max": [np.pi, np.pi, np.pi], + }, + } + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to", "normalization_modes", "target_rotations"} + else: + include = kwargs.pop("include", None) + + return super().model_dump(*args, include=include, **kwargs) + + @field_validator("modality_metadata", mode="before") + def validate_modality_metadata(cls, v): + for modality_key, config in v.items(): + if isinstance(config, dict): + config = StateActionMetadata.model_validate(config) + else: + assert isinstance( + config, StateActionMetadata + ), f"Invalid source rotation config: {config}" + v[modality_key] = config + return v + + @model_validator(mode="after") + def validate_normalization_statistics(self): + for modality_key, normalization_statistics in self.normalization_statistics.items(): + if modality_key in self.normalization_modes: + normalization_mode = self.normalization_modes[modality_key] + if normalization_mode == "min_max": + assert ( + "min" in normalization_statistics and "max" in normalization_statistics + ), f"Min and max statistics are required for min_max normalization, but got {normalization_statistics}" + assert len(normalization_statistics["min"]) == len( + normalization_statistics["max"] + ), f"Min and max statistics must have the same length, but got {normalization_statistics['min']} and {normalization_statistics['max']}" + elif normalization_mode == "mean_std": + assert ( + "mean" in normalization_statistics and "std" in normalization_statistics + ), f"Mean and std statistics are required for mean_std normalization, but got {normalization_statistics}" + assert len(normalization_statistics["mean"]) == len( + normalization_statistics["std"] + ), f"Mean and std statistics must have the same length, but got {normalization_statistics['mean']} and {normalization_statistics['std']}" + elif normalization_mode == "q99": + assert ( + "q01" in normalization_statistics and "q99" in normalization_statistics + ), f"q01 and q99 statistics are required for q99 normalization, but got {normalization_statistics}" + assert len(normalization_statistics["q01"]) == len( + normalization_statistics["q99"] + ), f"q01 and q99 statistics must have the same length, but got {normalization_statistics['q01']} and {normalization_statistics['q99']}" + elif normalization_mode == "binary": + assert ( + len(normalization_statistics) == 1 + ), f"Binary normalization should only have one value, but got {normalization_statistics}" + assert normalization_statistics[0] in [ + 0, + 1, + ], f"Binary normalization should only have 0 or 1, but got {normalization_statistics[0]}" + else: + raise ValueError(f"Invalid normalization mode: {normalization_mode}") + return self + + def set_metadata(self, dataset_metadata: DatasetMetadata): + dataset_statistics = dataset_metadata.statistics + modality_metadata = dataset_metadata.modalities + + # Check that all state keys specified in apply_to have their modality_metadata + for key in self.apply_to: + split_key = key.split(".", 1) + assert len(split_key) == 2, "State keys should have two parts: 'modality.key'" + if key not in self.modality_metadata: + modality, state_key = split_key + assert hasattr(modality_metadata, modality), f"{modality} config not found" + assert state_key in getattr( + modality_metadata, modality + ), f"{state_key} config not found" + self.modality_metadata[key] = getattr(modality_metadata, modality)[state_key] + + # Check that all state keys specified in normalization_modes have their statistics in state_statistics + for key in self.normalization_modes: + split_key = key.split(".", 1) + assert len(split_key) == 2, "State keys should have two parts: 'modality.key'" + modality, state_key = split_key + assert hasattr(dataset_statistics, modality), f"{modality} statistics not found" + assert state_key in getattr( + dataset_statistics, modality + ), f"{state_key} statistics not found" + assert ( + len(getattr(modality_metadata, modality)[state_key].shape) == 1 + ), f"{getattr(modality_metadata, modality)[state_key].shape=}" + self.normalization_statistics[key] = getattr(dataset_statistics, modality)[ + state_key + ].model_dump() + + # Initialize the rotation transformers + for key in self.target_rotations: + # Get the original representation of the state + from_rep = self.modality_metadata[key].rotation_type + assert from_rep is not None, f"Source rotation type not found for {key}" + + # Get the target representation of the state, will raise an error if the target representation is not valid + to_rep = RotationType(self.target_rotations[key]) + + # If the original representation is not the same as the target representation, initialize the rotation transformer + if from_rep != to_rep: + self._rotation_transformers[key] = RotationTransform( + from_rep=from_rep.value, to_rep=to_rep.value + ) + + # Initialize the normalizers + for key in self.normalization_modes: + modality, state_key = key.split(".", 1) + # If the state has a nontrivial rotation, we need to handle it more carefully + # For absolute rotations, we need to convert them to the target representation and normalize them using min_max mode, + # since we can infer the bounds by the representation + # For relative rotations, we cannot normalize them as we don't know the bounds + if key in self._rotation_transformers: + # Case 1: Absolute rotation + if self.modality_metadata[key].absolute: + # Check that the normalization mode is valid + assert ( + self.normalization_modes[key] == "min_max" + ), "Absolute rotations that are converted to other formats must be normalized using `min_max` mode" + rotation_type = RotationType(self.target_rotations[key]).value + # If the target representation is euler angles, we need to parse the convention + if rotation_type.startswith("euler_angles"): + rotation_type = "euler_angles" + # Get the statistics for the target representation + statistics = self._DEFAULT_MIN_MAX_STATISTICS[rotation_type] + # Case 2: Relative rotation + else: + raise ValueError( + f"Cannot normalize relative rotations: {key} that's converted to {self.target_rotations[key]}" + ) + # If the state is not continuous, we should not use normalization modes other than binary + elif ( + not self.modality_metadata[key].continuous + and self.normalization_modes[key] != "binary" + ): + raise ValueError( + f"{key} is not continuous, so it should be normalized using `binary` mode" + ) + # Initialize the normalizer + else: + statistics = self.normalization_statistics[key] + self._normalizers[key] = Normalizer( + mode=self.normalization_modes[key], statistics=statistics, + binary_threshold=self.binary_threshold, + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + # We allow some keys to be missing in the data, and only process the keys that are present + continue + if key not in self._input_dtypes: + input_dtype = data[key].dtype + assert isinstance( + input_dtype, torch.dtype + ), f"Unexpected input dtype: {input_dtype}. Expected type: {torch.dtype}" + self._input_dtypes[key] = input_dtype + else: + assert ( + data[key].dtype == self._input_dtypes[key] + ), f"All states corresponding to the same key must be of the same dtype, input dtype: {data[key].dtype}, expected dtype: {self._input_dtypes[key]}" + # Rotate the state + state = data[key] + if key in self._rotation_transformers: + state = self._rotation_transformers[key].forward(state) + # Normalize the state + if key in self._normalizers: + state = self._normalizers[key].forward(state) + data[key] = state + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + if key not in data: + continue + state = data[key] + assert isinstance( + state, torch.Tensor + ), f"Unexpected state type: {type(state)}. Expected type: {torch.Tensor}" + # Unnormalize the state + if key in self._normalizers: + state = self._normalizers[key].inverse(state) + # Change the state back to its original representation + if key in self._rotation_transformers: + state = self._rotation_transformers[key].inverse(state) + assert isinstance( + state, torch.Tensor + ), f"State should be tensor after unapplying transformations, but got {type(state)}" + # Only convert back to the original dtype if it's known, i.e. `apply` was called before + # If not, we don't know the original dtype, so we don't convert + if key in self._input_dtypes: + original_dtype = self._input_dtypes[key] + if isinstance(original_dtype, np.dtype): + state = state.numpy().astype(original_dtype) + elif isinstance(original_dtype, torch.dtype): + state = state.to(original_dtype) + else: + raise ValueError(f"Invalid input dtype: {original_dtype}") + data[key] = state + return data + + +class StateActionPerturbation(ModalityTransform): + """ + Class for state or action perturbation. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + std (float): Standard deviation of the noise to be added to the state or action. + """ + + # Configurable attributes + std: float = Field( + ..., description="Standard deviation of the noise to be added to the state or action." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Don't perturb the data in eval mode + return data + if self.std < 0: + # If the std is negative, we don't add any noise + return data + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + transformed_data_min = torch.min(state) + transformed_data_max = torch.max(state) + noise = torch.randn_like(state) * self.std + state += noise + # Clip to the original range + state = torch.clamp(state, transformed_data_min, transformed_data_max) + data[key] = state + return data + + +class StateActionDropout(ModalityTransform): + """ + Class for state or action dropout. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + dropout_prob (float): Probability of dropping out a state or action. + """ + + # Configurable attributes + dropout_prob: float = Field(..., description="Probability of dropping out a state or action.") + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Don't drop out the data in eval mode + return data + if self.dropout_prob < 0: + # If the dropout probability is negative, we don't drop out any states + return data + if self.dropout_prob > 1e-9 and random.random() < self.dropout_prob: + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + state = torch.zeros_like(state) + data[key] = state + return data + + +class StateActionSinCosTransform(ModalityTransform): + """ + Class for state or action sin-cos transform. + + Args: + apply_to (list[str]): The keys in the modality to load and transform. + """ + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + for key in self.apply_to: + state = data[key] + assert isinstance(state, torch.Tensor) + sin_state = torch.sin(state) + cos_state = torch.cos(state) + data[key] = torch.cat([sin_state, cos_state], dim=-1) + return data diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/video.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/video.py new file mode 100644 index 0000000000000000000000000000000000000000..15310f697259d7ee6ed8eac7e5abaca211b81dc9 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/transform/video.py @@ -0,0 +1,612 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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, Callable, ClassVar, Literal + +import albumentations as A +import cv2 +import numpy as np +import torch +import torchvision.transforms.v2 as T +from einops import rearrange +from pydantic import Field, PrivateAttr, field_validator +from PIL import Image + +from ..schema import DatasetMetadata +from .base import ModalityTransform + + +class VideoTransform(ModalityTransform): + # Configurable attributes + backend: str = Field( + default="torchvision", description="The backend to use for the transformations" + ) + + # Model variables + _train_transform: Callable | None = PrivateAttr(default=None) + _eval_transform: Callable | None = PrivateAttr(default=None) + _original_resolutions: dict[str, tuple[int, int]] = PrivateAttr(default_factory=dict) + + # Model constants + _INTERPOLATION_MAP: ClassVar[dict[str, dict[str, Any]]] = PrivateAttr( + { + "nearest": { + "albumentations": cv2.INTER_NEAREST, + "torchvision": T.InterpolationMode.NEAREST, + }, + "linear": { + "albumentations": cv2.INTER_LINEAR, + "torchvision": T.InterpolationMode.BILINEAR, + }, + "cubic": { + "albumentations": cv2.INTER_CUBIC, + "torchvision": T.InterpolationMode.BICUBIC, + }, + "area": { + "albumentations": cv2.INTER_AREA, + "torchvision": None, # Torchvision does not support this interpolation mode + }, + "lanczos4": { + "albumentations": cv2.INTER_LANCZOS4, # Lanczos with a 4x4 filter + "torchvision": T.InterpolationMode.LANCZOS, # Torchvision does not specify filter size, might be different from 4x4 + }, + "linear_exact": { + "albumentations": cv2.INTER_LINEAR_EXACT, + "torchvision": None, # Torchvision does not support this interpolation mode + }, + "nearest_exact": { + "albumentations": cv2.INTER_NEAREST_EXACT, + "torchvision": T.InterpolationMode.NEAREST_EXACT, + }, + "max": { + "albumentations": cv2.INTER_MAX, + "torchvision": None, + }, + } + ) + + @property + def train_transform(self) -> Callable: + assert ( + self._train_transform is not None + ), "Transform is not set. Please call set_metadata() before calling apply()." + return self._train_transform + + @train_transform.setter + def train_transform(self, value: Callable): + self._train_transform = value + + @property + def eval_transform(self) -> Callable | None: + return self._eval_transform + + @eval_transform.setter + def eval_transform(self, value: Callable | None): + self._eval_transform = value + + @property + def original_resolutions(self) -> dict[str, tuple[int, int]]: + assert ( + self._original_resolutions is not None + ), "Original resolutions are not set. Please call set_metadata() before calling apply()." + return self._original_resolutions + + @original_resolutions.setter + def original_resolutions(self, value: dict[str, tuple[int, int]]): + self._original_resolutions = value + + def check_input(self, data: dict[str, Any]): + if self.backend == "torchvision": + for key in self.apply_to: + assert isinstance(data[key], torch.Tensor), f"Video {key} is not a torch tensor" + assert data[key].ndim in [ + 4, + 5, + ], f"Expected video {key} to have 4 or 5 dimensions (T, C, H, W or T, B, C, H, W), got {data[key].ndim}" + elif self.backend == "albumentations": + for key in self.apply_to: + assert isinstance(data[key], np.ndarray), f"Video {key} is not a numpy array" + assert data[key].ndim in [ + 4, + 5, + ], f"Expected video {key} to have 4 or 5 dimensions (T, C, H, W or T, B, C, H, W), got {data[key].ndim}" + else: + raise ValueError(f"Backend {self.backend} not supported") + + def set_metadata(self, dataset_metadata: DatasetMetadata): + super().set_metadata(dataset_metadata) + self.original_resolutions = {} + for key in self.apply_to: + split_keys = key.split(".") + assert len(split_keys) == 2, f"Invalid key: {key}. Expected format: modality.key" + sub_key = split_keys[1] + if sub_key in dataset_metadata.modalities.video: + self.original_resolutions[key] = dataset_metadata.modalities.video[ + sub_key + ].resolution + else: + raise ValueError( + f"Video key {sub_key} not found in dataset metadata. Available keys: {dataset_metadata.modalities.video.keys()}" + ) + train_transform = self.get_transform(mode="train") + eval_transform = self.get_transform(mode="eval") + if self.backend == "albumentations": + self.train_transform = A.ReplayCompose(transforms=[train_transform]) # type: ignore + if eval_transform is not None: + self.eval_transform = A.ReplayCompose(transforms=[eval_transform]) # type: ignore + else: + assert train_transform is not None, "Train transform must be set" + self.train_transform = train_transform + self.eval_transform = eval_transform + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if self.training: + transform = self.train_transform + else: + transform = self.eval_transform + if transform is None: + return data + assert ( + transform is not None + ), "Transform is not set. Please call set_metadata() before calling apply()." + try: + self.check_input(data) + except AssertionError as e: + raise ValueError( + f"Input data does not match the expected format for {self.__class__.__name__}: {e}" + ) from e + + # Concatenate views + views = [data[key] for key in self.apply_to] + num_views = len(views) + is_batched = views[0].ndim == 5 + bs = views[0].shape[0] if is_batched else 1 + if isinstance(views[0], torch.Tensor): + views = torch.cat(views, 0) + elif isinstance(views[0], np.ndarray): + views = np.concatenate(views, 0) + else: + raise ValueError(f"Unsupported view type: {type(views[0])}") + if is_batched: + views = rearrange(views, "(v b) t c h w -> (v b t) c h w", v=num_views, b=bs) + # Apply the transform + if self.backend == "torchvision": + views = transform(views) + elif self.backend == "albumentations": + assert isinstance(transform, A.ReplayCompose), "Transform must be a ReplayCompose" + first_frame = views[0] + transformed = transform(image=first_frame) + replay_data = transformed["replay"] + transformed_first_frame = transformed["image"] + + if len(views) > 1: + # Apply the same transformations to the rest of the frames + transformed_frames = [ + transform.replay(replay_data, image=frame)["image"] for frame in views[1:] + ] + # Add the first frame back + transformed_frames = [transformed_first_frame] + transformed_frames + else: + # If there is only one frame, just make a list with one frame + transformed_frames = [transformed_first_frame] + + # Delete the replay data to save memory + del replay_data + views = np.stack(transformed_frames, 0) + + else: + raise ValueError(f"Backend {self.backend} not supported") + # Split views + if is_batched: + views = rearrange(views, "(v b t) c h w -> v b t c h w", v=num_views, b=bs) + else: + views = rearrange(views, "(v t) c h w -> v t c h w", v=num_views) + for key, view in zip(self.apply_to, views): + data[key] = view + return data + + @classmethod + def _validate_interpolation(cls, interpolation: str): + if interpolation not in cls._INTERPOLATION_MAP: + raise ValueError(f"Interpolation mode {interpolation} not supported") + + def _get_interpolation(self, interpolation: str, backend: str = "torchvision"): + """ + Get the interpolation mode for the given backend. + + Args: + interpolation (str): The interpolation mode. + backend (str): The backend to use. + + Returns: + Any: The interpolation mode for the given backend. + """ + return self._INTERPOLATION_MAP[interpolation][backend] + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + raise NotImplementedError( + "set_transform is not implemented for VideoTransform. Please implement this function to set the transforms." + ) + + +class VideoCrop(VideoTransform): + height: int | None = Field(default=None, description="The height of the input image") + width: int | None = Field(default=None, description="The width of the input image") + scale: float = Field( + ..., + description="The scale of the crop. The crop size is (width * scale, height * scale)", + ) + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the transform for the given mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: If mode is "train", return a random crop transform. If mode is "eval", return a center crop transform. + """ + # 1. Check the input resolution + assert ( + len(set(self.original_resolutions.values())) == 1 + ), f"All video keys must have the same resolution, got: {self.original_resolutions}" + if self.height is None: + assert self.width is None, "Height and width must be either both provided or both None" + self.width, self.height = self.original_resolutions[self.apply_to[0]] + else: + assert ( + self.width is not None + ), "Height and width must be either both provided or both None" + # 2. Create the transform + size = (int(self.height * self.scale), int(self.width * self.scale)) + if self.backend == "torchvision": + if mode == "train": + return T.RandomCrop(size) + elif mode == "eval": + return T.CenterCrop(size) + else: + raise ValueError(f"Crop mode {mode} not supported") + elif self.backend == "albumentations": + if mode == "train": + return A.RandomCrop(height=size[0], width=size[1], p=1) + elif mode == "eval": + return A.CenterCrop(height=size[0], width=size[1], p=1) + else: + raise ValueError(f"Crop mode {mode} not supported") + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict[str, Any]): + super().check_input(data) + # Check the input resolution + for key in self.apply_to: + if self.backend == "torchvision": + height, width = data[key].shape[-2:] + elif self.backend == "albumentations": + height, width = data[key].shape[-3:-1] + else: + raise ValueError(f"Backend {self.backend} not supported") + assert ( + height == self.height and width == self.width + ), f"Video {key} has invalid shape {height, width}, expected {self.height, self.width}" + + +class VideoResize(VideoTransform): + height: int = Field(..., description="The height of the resize") + width: int = Field(..., description="The width of the resize") + interpolation: str = Field(default="linear", description="The interpolation mode") + antialias: bool = Field(default=True, description="Whether to apply antialiasing") + + @field_validator("interpolation") + def validate_interpolation(cls, v): + cls._validate_interpolation(v) + return v + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the resize transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The resize transform. + """ + interpolation = self._get_interpolation(self.interpolation, self.backend) + if interpolation is None: + raise ValueError( + f"Interpolation mode {self.interpolation} not supported for torchvision" + ) + if self.backend == "torchvision": + size = (self.height, self.width) + return T.Resize(size, interpolation=interpolation, antialias=self.antialias) + elif self.backend == "albumentations": + return A.Resize( + height=self.height, + width=self.width, + interpolation=interpolation, + p=1, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomRotation(VideoTransform): + degrees: float | tuple[float, float] = Field( + ..., description="The degrees of the random rotation" + ) + interpolation: str = Field("linear", description="The interpolation mode") + + @field_validator("interpolation") + def validate_interpolation(cls, v): + cls._validate_interpolation(v) + return v + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the random rotation transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: The random rotation transform. None for eval mode. + """ + if mode == "eval": + return None + interpolation = self._get_interpolation(self.interpolation, self.backend) + if interpolation is None: + raise ValueError( + f"Interpolation mode {self.interpolation} not supported for torchvision" + ) + if self.backend == "torchvision": + return T.RandomRotation(self.degrees, interpolation=interpolation) # type: ignore + elif self.backend == "albumentations": + return A.Rotate(limit=self.degrees, interpolation=interpolation, p=1) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoHorizontalFlip(VideoTransform): + p: float = Field(..., description="The probability of the horizontal flip") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the horizontal flip transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a horizontal flip transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomHorizontalFlip(self.p) + elif self.backend == "albumentations": + return A.HorizontalFlip(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoGrayscale(VideoTransform): + p: float = Field(..., description="The probability of the grayscale transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the grayscale transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a grayscale transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomGrayscale(self.p) + elif self.backend == "albumentations": + return A.ToGray(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoColorJitter(VideoTransform): + brightness: float | tuple[float, float] = Field( + ..., description="The brightness of the color jitter" + ) + contrast: float | tuple[float, float] = Field( + ..., description="The contrast of the color jitter" + ) + saturation: float | tuple[float, float] = Field( + ..., description="The saturation of the color jitter" + ) + hue: float | tuple[float, float] = Field(..., description="The hue of the color jitter") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the color jitter transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a color jitter transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.ColorJitter( + brightness=self.brightness, + contrast=self.contrast, + saturation=self.saturation, + hue=self.hue, + ) + elif self.backend == "albumentations": + return A.ColorJitter( + brightness=self.brightness, + contrast=self.contrast, + saturation=self.saturation, + hue=self.hue, + p=1, + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomGrayscale(VideoTransform): + p: float = Field(..., description="The probability of the grayscale transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the grayscale transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a grayscale transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomGrayscale(self.p) + elif self.backend == "albumentations": + return A.ToGray(p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoRandomPosterize(VideoTransform): + bits: int = Field(..., description="The number of bits to posterize the image") + p: float = Field(..., description="The probability of the posterize transformation") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """Get the posterize transform, only used in train mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable | None: If mode is "train", return a posterize transform. If mode is "eval", return None. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomPosterize(bits=self.bits, p=self.p) + elif self.backend == "albumentations": + return A.Posterize(num_bits=self.bits, p=self.p) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +class VideoToTensor(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to tensor transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to tensor transform. + """ + if self.backend == "torchvision": + return self.__class__.to_tensor + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict): + """Check if the input data has the correct shape. + Expected video shape: [T, H, W, C], dtype np.uint8 + """ + for key in self.apply_to: + assert key in data, f"Key {key} not found in data. Available keys: {data.keys()}" + assert data[key].ndim in [ + 4, + 5, + ], f"Video {key} must have 4 or 5 dimensions, got {data[key].ndim}" + assert ( + data[key].dtype == np.uint8 + ), f"Video {key} must have dtype uint8, got {data[key].dtype}" + input_resolution = data[key].shape[-3:-1][::-1] + if key in self.original_resolutions: + expected_resolution = self.original_resolutions[key] + else: + expected_resolution = input_resolution + assert ( + input_resolution == expected_resolution + ), f"Video {key} has invalid resolution {input_resolution}, expected {expected_resolution}. Full shape: {data[key].shape}" + + @staticmethod + def to_tensor(frames: np.ndarray) -> torch.Tensor: + """Convert numpy array to tensor efficiently. + + Args: + frames: numpy array of shape [T, H, W, C] in uint8 format + Returns: + tensor of shape [T, C, H, W] in range [0, 1] + """ + frames_tensor = torch.from_numpy(frames).to(torch.float32) / 255.0 + return frames_tensor.permute(0, 3, 1, 2) # [T, C, H, W] + + +class VideoToNumpy(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to numpy transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to numpy transform. + """ + if self.backend == "torchvision": + return self.__class__.to_numpy + else: + raise ValueError(f"Backend {self.backend} not supported") + + @staticmethod + def to_numpy(frames: torch.Tensor) -> np.ndarray: + """Convert tensor back to numpy array efficiently. + + Args: + frames: tensor of shape [T, C, H, W] in range [0, 1] + Returns: + numpy array of shape [T, H, W, C] in uint8 format + """ + return (frames.permute(0, 2, 3, 1) * 255).to(torch.uint8).cpu().numpy() + +class VideoToPIL(VideoTransform): + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the to PIL transform. Same transform for both train and eval. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The to PIL transform. + """ + if self.backend == "torchvision": + return self.__class__.to_pil + else: + raise ValueError(f"Backend {self.backend} not supported") + + @staticmethod + def to_pil(frames: torch.Tensor) -> Image.Image: + """Convert tensor back to PIL Image. + + Args: + frames: tensor of shape [T, C, H, W] in range [0, 1] + Returns: + PIL Image of shape [T, H, W, C] in uint8 format + """ + # video PIL format? + return Image.fromarray((frames.permute(0, 2, 3, 1) * 255).to(torch.uint8).cpu().numpy()) \ No newline at end of file diff --git a/real_robot/source/starVLA/dataloader/gr00t_lerobot/video.py b/real_robot/source/starVLA/dataloader/gr00t_lerobot/video.py new file mode 100644 index 0000000000000000000000000000000000000000..6d61e1c2f53f0c5c4b5c751923922a3df684a3ca --- /dev/null +++ b/real_robot/source/starVLA/dataloader/gr00t_lerobot/video.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + + +import av +import cv2 +import numpy as np + +import torch # noqa: F401 # isort: skip +import torchvision # noqa: F401 # isort: skip + +# Import decord with graceful fallback +try: + import decord # noqa: F401 + + DECORD_AVAILABLE = True +except ImportError: + DECORD_AVAILABLE = False + +try: + import torchcodec + + TORCHCODEC_AVAILABLE = True +except (ImportError, RuntimeError): + TORCHCODEC_AVAILABLE = False + + +def get_frames_by_indices( + video_path: str, + indices: list[int] | np.ndarray, + video_backend: str = "decord", + video_backend_kwargs: dict = {}, +) -> np.ndarray: + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available.") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + frames = vr.get_batch(indices) + return frames.asnumpy() + elif video_backend == "torchcodec": + if not TORCHCODEC_AVAILABLE: + raise ImportError("torchcodec is not available.") + decoder = torchcodec.decoders.VideoDecoder( + video_path, device="cpu", dimension_order="NHWC", num_ffmpeg_threads=0 + ) + return decoder.get_frames_at(indices=indices).data.numpy() + elif video_backend == "opencv": + frames = [] + cap = cv2.VideoCapture(video_path, **video_backend_kwargs) + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if not ret: + raise ValueError(f"Unable to read frame at index {idx}") + frames.append(frame) + cap.release() + frames = np.array(frames) + return frames + elif video_backend == "pyav": + # Use PyAV directly without torchvision wrapper + container = None + try: + container = av.open(video_path) + stream = container.streams.video[0] + time_base = float(stream.time_base) + fps = float(stream.average_rate) if stream.average_rate else float(stream.guessed_rate) + + loaded_frames = [] + + for target_idx in indices: + # Convert frame index to timestamp, then to pts + target_ts = target_idx / fps + target_pts = int(target_ts / time_base) + + # Seek to the target position + container.seek(target_pts, stream=stream, backward=True, any_frame=False) + + closest_frame = None + closest_idx_diff = float('inf') + + for frame in container.decode(video=0): + current_ts = float(frame.pts * time_base) + current_idx = int(round(current_ts * fps)) + current_diff = abs(current_idx - target_idx) + + if current_diff < closest_idx_diff: + closest_idx_diff = current_diff + closest_frame = frame + + # If we've passed the target and diff is increasing, stop + if current_idx > target_idx and current_diff > closest_idx_diff: + break + + # Stop if we're significantly past the target + if current_idx > target_idx + 5: + break + + if closest_frame is not None: + frame_data = closest_frame.to_ndarray(format="rgb24") + loaded_frames.append(frame_data) + else: + raise ValueError(f"Unable to find frame at index {target_idx}") + + frames = np.array(loaded_frames) + return frames + + finally: + if container is not None: + container.close() + else: + raise NotImplementedError + + +def get_frames_by_timestamps( + video_path: str, + timestamps: list[float] | np.ndarray, + video_backend: str = "decord", + video_backend_kwargs: dict = {}, +) -> np.ndarray: + """Get frames from a video at specified timestamps. + Args: + video_path (str): Path to the video file. + timestamps (list[int] | np.ndarray): Timestamps to retrieve frames for, in seconds. + video_backend (str, optional): Video backend to use. Defaults to "decord". + Returns: + np.ndarray: Frames at the specified timestamps. + """ + if video_backend == "decord": + # For some GPUs, AV format data cannot be read + if not DECORD_AVAILABLE: + raise ImportError("decord is not available.") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + num_frames = len(vr) + # Retrieve the timestamps for each frame in the video + frame_ts: np.ndarray = vr.get_frame_timestamp(range(num_frames)) + # Map each requested timestamp to the closest frame index + # Only take the first element of the frame_ts array which corresponds to start_seconds + indices = np.abs(frame_ts[:, :1] - timestamps).argmin(axis=0) + frames = vr.get_batch(indices) + return frames.asnumpy() + elif video_backend == "torchcodec": + if not TORCHCODEC_AVAILABLE: + raise ImportError("torchcodec is not available.") + decoder = torchcodec.decoders.VideoDecoder( + video_path, device="cpu", dimension_order="NHWC", num_ffmpeg_threads=0 + ) + return decoder.get_frames_played_at(seconds=timestamps).data.numpy() + elif video_backend == "opencv": + # Open the video file + cap = cv2.VideoCapture(video_path, **video_backend_kwargs) + if not cap.isOpened(): + raise ValueError(f"Unable to open video file: {video_path}") + # Retrieve the total number of frames + num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + # Calculate timestamps for each frame + fps = cap.get(cv2.CAP_PROP_FPS) + frame_ts = np.arange(num_frames) / fps + frame_ts = frame_ts[:, np.newaxis] # Reshape to (num_frames, 1) for broadcasting + # Map each requested timestamp to the closest frame index + indices = np.abs(frame_ts - timestamps).argmin(axis=0) + frames = [] + for idx in indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, frame = cap.read() + if not ret: + raise ValueError(f"Unable to read frame at index {idx}") + frames.append(frame) + cap.release() + frames = np.array(frames) + return frames + elif video_backend == "pyav": + # Use PyAV directly without torchvision wrapper + container = None + try: + container = av.open(video_path) + stream = container.streams.video[0] + + # Get video properties + time_base = float(stream.time_base) + fps = float(stream.average_rate) if stream.average_rate else float(stream.guessed_rate) + duration = float(stream.duration * time_base) if stream.duration else None + + loaded_frames = [] + + for target_ts in timestamps: + # Convert timestamp to pts (presentation timestamp) + target_pts = int(target_ts / time_base) + + # Seek to the target timestamp (seek to keyframe before target) + container.seek(target_pts, stream=stream, backward=True, any_frame=False) + + closest_frame = None + closest_ts_diff = float('inf') + + for frame in container.decode(video=0): + current_ts = float(frame.pts * time_base) + current_diff = abs(current_ts - target_ts) + + if current_diff < closest_ts_diff: + # Release the previous frame + if closest_frame is not None: + del closest_frame + closest_ts_diff = current_diff + closest_frame = frame + + # If we've passed the target and diff is increasing, stop + if current_ts > target_ts and current_diff > closest_ts_diff: + break + + # Also stop if we're significantly past the target (optimization) + if current_ts > target_ts + 1.0: + break + + if closest_frame is not None: + frame_data = closest_frame.to_ndarray(format="rgb24") + loaded_frames.append(frame_data) + del closest_frame + else: + raise ValueError(f"Unable to find frame at timestamp {target_ts}") + + frames = np.array(loaded_frames) + return frames + + finally: + if container is not None: + container.close() + container = None + + elif video_backend == "torchvision_av": + torchvision.set_video_backend("pyav") + loaded_frames = [] + loaded_ts = [] + + reader = None + try: + reader = torchvision.io.VideoReader(video_path, "video") + + for target_ts in timestamps: + # Reset reader state + reader.seek(target_ts, keyframes_only=True) + + closest_frame = None + closest_ts_diff = float('inf') + + for frame in reader: + current_ts = frame["pts"] + current_diff = abs(current_ts - target_ts) + + if closest_frame is None: + closest_frame = frame + + if current_diff < closest_ts_diff: + # Release the previous frame + if closest_frame is not None: + del closest_frame + closest_ts_diff = current_diff + closest_frame = frame + else: + # The time difference starts to increase, stop searching + break + + if closest_frame is not None: + frame_data = closest_frame["data"] + if isinstance(frame_data, torch.Tensor): + frame_data = frame_data.cpu().numpy() + loaded_frames.append(frame_data) + loaded_ts.append(closest_frame["pts"]) + + # Immediately release frame reference + del closest_frame + + finally: + # Thoroughly clean resources + if reader is not None: + if hasattr(reader, '_c'): + reader._c = None + if hasattr(reader, 'container'): + reader.container.close() + reader.container = None + + frames = np.array(loaded_frames) + return frames.transpose(0, 2, 3, 1) + else: + raise NotImplementedError + + +def get_all_frames( + video_path: str, + video_backend: str = "decord", + video_backend_kwargs: dict = {}, + resize_size: tuple[int, int] | None = None, +) -> np.ndarray: + """Get all frames from a video. + Args: + video_path (str): Path to the video file. + video_backend (str, optional): Video backend to use. Defaults to "decord". + video_backend_kwargs (dict, optional): Keyword arguments for the video backend. + resize_size (tuple[int, int], optional): Resize size for the frames. Defaults to None. + """ + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available.") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + frames = vr.get_batch(range(len(vr))).asnumpy() + elif video_backend == "torchcodec": + if not TORCHCODEC_AVAILABLE: + raise ImportError("torchcodec is not available.") + decoder = torchcodec.decoders.VideoDecoder( + video_path, device="cpu", dimension_order="NHWC", num_ffmpeg_threads=0 + ) + frames = decoder.get_frames_at(indices=range(len(decoder))) + return frames.data.numpy(), frames.pts_seconds.numpy() + elif video_backend == "pyav": + container = av.open(video_path) + frames = [] + for frame in container.decode(video=0): + frame = frame.to_ndarray(format="rgb24") + frames.append(frame) + frames = np.array(frames) + elif video_backend == "torchvision_av": + # set backend and reader + torchvision.set_video_backend("pyav") + reader = torchvision.io.VideoReader(video_path, "video") + frames = [] + for frame in reader: + frames.append(frame["data"].numpy()) + frames = np.array(frames) + frames = frames.transpose(0, 2, 3, 1) + else: + raise NotImplementedError(f"Video backend {video_backend} not implemented") + # resize frames if specified + if resize_size is not None: + frames = [cv2.resize(frame, resize_size) for frame in frames] + frames = np.array(frames) + return frames \ No newline at end of file diff --git a/real_robot/source/starVLA/dataloader/infoflow_sampler.py b/real_robot/source/starVLA/dataloader/infoflow_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..db228ababf79814d5c12e3dc615cffe4b13ea9cf --- /dev/null +++ b/real_robot/source/starVLA/dataloader/infoflow_sampler.py @@ -0,0 +1,6 @@ +"""Compatibility alias for the ReTime plugin; preserves legacy imports.""" + +from importlib import import_module +import sys + +sys.modules[__name__] = import_module('retime.sampler') diff --git a/real_robot/source/starVLA/dataloader/lerobot_datasets.py b/real_robot/source/starVLA/dataloader/lerobot_datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..8c4134eb77d159418979fd5f0f59f073f59c9dd9 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/lerobot_datasets.py @@ -0,0 +1,162 @@ +# Modified in the ReTime research artifact; see NOTICE. +# Copyright 2025 NVIDIA Corp. and affiliates. All rights reserved. +# Modification: [return raw data and suport multi-dataset mixture]. +# Modification: [suport topdowm processing, suport param from config]. + +import logging +from pathlib import Path +from typing import Sequence +from omegaconf import OmegaConf + +from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset, LeRobotMixtureDataset +from starVLA.dataloader.gr00t_lerobot.registry import ( + ROBOT_TYPE_CONFIG_MAP, + DATASET_NAMED_MIXTURES, + EmbodimentTag, +) + +logger = logging.getLogger(__name__) + +def collate_fn(batch): + return batch + +def make_LeRobotSingleDataset( + data_root_dir: Path | str, + data_name: str, + robot_type: str, + delete_pause_frame: bool = False, + data_cfg: dict | None = None, +) -> LeRobotSingleDataset: + """ + Make a LeRobotSingleDataset object. + + :param data_root_dir: The root directory of the dataset. + :param data_name: The name of the dataset. + :param robot_type: The robot type config to use. + :param crop_obs_camera: Whether to crop the observation camera images. + :return: A LeRobotSingleDataset object. + """ + + data_config = ROBOT_TYPE_CONFIG_MAP[robot_type] + modality_config = data_config.modality_config() + transforms = data_config.transform() + dataset_path = data_root_dir / data_name + embodiment_tag = getattr(data_config, "embodiment_tag", None) + if embodiment_tag is None: + print(f"Warning: DataConfig for robot_type={robot_type!r} has no embodiment_tag, using {EmbodimentTag.NEW_EMBODIMENT} as default") + embodiment_tag = EmbodimentTag.NEW_EMBODIMENT + + video_backend = data_cfg.get("video_backend", "decord") if data_cfg else "torchvision_av" + + # Opt-in factory hook: a DataConfig may define ``make_dataset(dataset_name=..., **ds_kwargs)`` + # to swap in a custom dataset class (e.g. with per-task filtering / chunk stride). + # When absent, fall through to the default LeRobotSingleDataset construction below. + if hasattr(data_config, "make_dataset"): + return data_config.make_dataset( + dataset_path=dataset_path, + modality_configs=modality_config, + transforms=transforms, + embodiment_tag=embodiment_tag, + video_backend=video_backend, + delete_pause_frame=delete_pause_frame, + data_cfg=data_cfg, + dataset_name=data_name, + ) + + return LeRobotSingleDataset( + dataset_path=dataset_path, + modality_configs=modality_config, + transforms=transforms, + embodiment_tag=embodiment_tag, + video_backend=video_backend, # decord is more efficiency | torchvision_av for video.av1 + delete_pause_frame=delete_pause_frame, + data_cfg=data_cfg, + dataset_name=data_name, + ) + +def get_vla_dataset( + data_cfg: dict, + mode: str = "train", + balance_dataset_weights: bool = False, + balance_trajectory_weights: bool = False, + seed: int = 42, + **kwargs: dict, +) -> LeRobotMixtureDataset: + """ + Get a LeRobotMixtureDataset object. + """ + data_root_dir = data_cfg.data_root_dir + data_mix = data_cfg.data_mix + delete_pause_frame = data_cfg.get("delete_pause_frame", False) + mixture_spec = DATASET_NAMED_MIXTURES[data_mix] + logger.info(f"[dataloader] Using mixture '{data_mix}': {[(d, w, r) for d, w, r in mixture_spec]}") + included_datasets, filtered_mixture_spec = set(), [] + for d_name, d_weight, robot_type in mixture_spec: + dataset_key = (d_name, robot_type) + if dataset_key in included_datasets: + print(f"Skipping Duplicate Dataset: `{(d_name, d_weight, robot_type)}`") + continue + + included_datasets.add(dataset_key) + filtered_mixture_spec.append((d_name, d_weight, robot_type)) + + dataset_mixture = [] + for d_name, d_weight, robot_type in filtered_mixture_spec: + dataset_mixture.append((make_LeRobotSingleDataset(Path(data_root_dir), d_name, robot_type, delete_pause_frame=delete_pause_frame, data_cfg=data_cfg), d_weight)) + + return LeRobotMixtureDataset( + dataset_mixture, + mode=mode, + balance_dataset_weights=balance_dataset_weights, + balance_trajectory_weights=balance_trajectory_weights, + seed=seed, + data_cfg=data_cfg, + **kwargs, + ) + + + +if __name__ == "__main__": + import argparse + import os + + parser = argparse.ArgumentParser() + parser.add_argument("--config_yaml", type=str, default="./examples/LIBERO/train_files/bar/starvla_cotrain_libero.yaml", help="Path to YAML config") + parser.add_argument("--data_mix", type=str, default=None, help="Override data_mix from config") + parser.add_argument("--data_root_dir", type=str, default=None, help="Override data_root_dir from config") + args = parser.parse_args() + + if os.getenv("DEBUGPY_ENABLE", "0") == "1": + import debugpy + debugpy.listen(("0.0.0.0", 10092)) + print("Rank 0 waiting for debugger attach on port 10092...") + debugpy.wait_for_client() + + cfg = OmegaConf.load(args.config_yaml) + vla_dataset_cfg = cfg.datasets.vla_data + vla_dataset_cfg.data_root_dir = Path(vla_dataset_cfg.data_root_dir) + if args.data_mix is not None: + vla_dataset_cfg.data_mix = args.data_mix + if args.data_root_dir is not None: + vla_dataset_cfg.data_root_dir = Path(args.data_root_dir) + + dataset = get_vla_dataset(data_cfg=vla_dataset_cfg) + from torch.utils.data import DataLoader + train_dataloader = DataLoader( + dataset, + batch_size=2, + num_workers=1, # For Debug + collate_fn=collate_fn, + ) + + cfg.output_dir = "./results/debug" + output_dir = Path(cfg.output_dir) + dataset.save_dataset_statistics(output_dir / "dataset_statistics.json") + + from tqdm import tqdm + count = 0 + for batch in tqdm(train_dataloader, desc="Processing Batches"): + if count > 3: + break + count += 1 + pass diff --git a/real_robot/source/starVLA/dataloader/relative_pose.py b/real_robot/source/starVLA/dataloader/relative_pose.py new file mode 100644 index 0000000000000000000000000000000000000000..943c9e15c1ac38fc894586a7b333d0e9c1c2f06a --- /dev/null +++ b/real_robot/source/starVLA/dataloader/relative_pose.py @@ -0,0 +1,96 @@ +"""SE(3) relative pose helpers for wxyz quaternions. + +``action_mode: rel`` must not subtract quaternion components. Positions are +``xyz_a - xyz_s0``; rotations are ``q_rel = inv(q_s0) ⊗ q_a``. The relative +quaternion is still a unit quaternion, so later rotation-6D conversion stays +valid with ``absolute: true`` in modality metadata. +""" + +from __future__ import annotations + +import numpy as np + +_QUATERNION_DIM = 4 + + +def _as_float(values: np.ndarray) -> np.ndarray: + return np.asarray(values, dtype=np.float32) + + +def quat_conjugate(quaternion: np.ndarray) -> np.ndarray: + """Inverse of a unit wxyz quaternion (conjugate).""" + conjugated = np.array(_as_float(quaternion), copy=True) + conjugated[..., 1:] *= -1 + return conjugated + + +def quat_multiply(left: np.ndarray, right: np.ndarray) -> np.ndarray: + """Hamilton product of wxyz quaternions, broadcasting on leading axes.""" + left_q, right_q = np.broadcast_arrays(_as_float(left), _as_float(right)) + w1, x1, y1, z1 = left_q[..., 0], left_q[..., 1], left_q[..., 2], left_q[..., 3] + w2, x2, y2, z2 = right_q[..., 0], right_q[..., 1], right_q[..., 2], right_q[..., 3] + product = np.stack( + [ + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ], + axis=-1, + ) + return _normalize_quaternion(product) + + +def _normalize_quaternion(quaternion: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(quaternion, axis=-1, keepdims=True) + return quaternion / np.clip(norms, 1e-8, None) + + +def relative_quaternion(state_quaternion: np.ndarray, action_quaternion: np.ndarray) -> np.ndarray: + """Rotation taking ``state`` onto ``action``: ``inv(q_s) ⊗ q_a``.""" + return quat_multiply(quat_conjugate(state_quaternion), action_quaternion) + + +def compose_quaternion(state_quaternion: np.ndarray, relative_quaternion_value: np.ndarray) -> np.ndarray: + """Apply a relative rotation: ``q_s ⊗ q_rel``.""" + return quat_multiply(state_quaternion, relative_quaternion_value) + + +def relative_values( + action_values: np.ndarray, + state0: np.ndarray, + *, + is_quaternion: bool, +) -> np.ndarray: + """Convert one action block into the ``rel`` representation.""" + action = _as_float(action_values) + origin = _as_float(state0) + if is_quaternion: + if action.shape[-1] != _QUATERNION_DIM or origin.shape[-1] != _QUATERNION_DIM: + raise ValueError( + f"quaternion rel requires last dim {_QUATERNION_DIM}, " + f"got action {action.shape} state {origin.shape}" + ) + return relative_quaternion(origin, action) + return action - origin + + +def compose_values( + relative_action: np.ndarray, + state0: np.ndarray, + *, + is_quaternion: bool, +) -> np.ndarray: + """Invert :func:`relative_values`.""" + relative = _as_float(relative_action) + origin = _as_float(state0) + if is_quaternion: + return compose_quaternion(origin, relative) + return relative + origin + + +def rotation_type_is_quaternion(rotation_type: object | None) -> bool: + if rotation_type is None: + return False + value = rotation_type.value if hasattr(rotation_type, "value") else str(rotation_type) + return value == "quaternion" diff --git a/real_robot/source/starVLA/dataloader/temporal_proposal.py b/real_robot/source/starVLA/dataloader/temporal_proposal.py new file mode 100644 index 0000000000000000000000000000000000000000..777256fb7488728579facb8a367d7aed499ef834 --- /dev/null +++ b/real_robot/source/starVLA/dataloader/temporal_proposal.py @@ -0,0 +1,6 @@ +"""Compatibility alias for the ReTime plugin; preserves legacy imports.""" + +from importlib import import_module +import sys + +sys.modules[__name__] = import_module('retime.proposal') diff --git a/real_robot/source/starVLA/model/framework/VLM4A/QwenGR00T.py b/real_robot/source/starVLA/model/framework/VLM4A/QwenGR00T.py new file mode 100644 index 0000000000000000000000000000000000000000..0923e4c87874bae484c1e62f3cb1b8cbe0c5264c --- /dev/null +++ b/real_robot/source/starVLA/model/framework/VLM4A/QwenGR00T.py @@ -0,0 +1,367 @@ +# Modified in the ReTime research artifact; see NOTICE. +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); +""" +Qwen-GR00T Framework +A lightweight implementation that Qwen-VL + Flow-matching head to directly predict continuous actions +Flow-matching header is copyright from GR00T N1.5, +""" + +import sys +from pathlib import Path + +# Add workspace root to Python path if not already there +_workspace_root = Path(__file__).parent.parent.parent.parent.parent +if str(_workspace_root) not in sys.path: + sys.path.insert(0, str(_workspace_root)) + +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +import numpy as np +import torch +from PIL import Image + +from deployment.model_server.tools.image_tools import to_pil_preserve +from starVLA.training.infoflow import aggregate_repeated_losses +from starVLA.training.trainer_utils import initialize_overwatch + +logger = initialize_overwatch(__name__) + +# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels) +IGNORE_INDEX = -100 + + +from starVLA.model.framework.base_framework import baseframework +from starVLA.model.framework.share_tools import merge_framework_config +from starVLA.model.modules.action_model.GR00T_ActionHeader import FlowmatchingActionHead, get_action_model +from starVLA.model.modules.vlm import get_vlm_model +from starVLA.model.tools import FRAMEWORK_REGISTRY +from starVLA.training.trainer_utils.trainer_tools import resize_images + + +# ────────────────────────────────────────────────────────────────────── +# Default Config for QwenGR00T +# - Documents every framework-level parameter with type + description +# - YAML values override these defaults; extra YAML keys are preserved +# ────────────────────────────────────────────────────────────────────── +@dataclass +class QwenGR00TDefaultConfig: + """QwenGR00T framework default parameters. + + All fields can be overridden by the corresponding key in the YAML + ``framework:`` section. Extra YAML keys not listed here are kept + as-is (Config-as-API flexibility). + """ + + # --- Registry identifier --- + name: str = "QwenGR00T" + + # === VLM backbone (Qwen2.5-VL / Qwen3-VL) === + qwenvl: dict = field( + default_factory=lambda: { + # Path to base VLM checkpoint (local or HF hub id) + "base_vlm": "./playground/Pretrained_models/Qwen3-VL-4B-Instruct", + # Attention implementation: "flash_attention_2" | "eager" | "sdpa" + "attn_implementation": "flash_attention_2", + # VLM hidden dimension (used for cross-attention alignment) + "vl_hidden_dim": 2048, + } + ) + + # # === DINO encoder (optional multi-view spatial tokens) === Dino is not used in this QwenGR00T version, we can add it later when we want to use it + # dino: dict = field(default_factory=lambda: { + # # DINO backbone variant: "dinov2_vits14" | "dinov2_vitb14" | ... + # "dino_backbone": "dinov2_vits14", + # }) + + # === Action head (Flow-matching / DiT diffusion) === + action_model: dict = field( + default_factory=lambda: { + # DiT model size: "DiT-B" | "DiT-L" | "DiT-XL" + "action_model_type": "DiT-B", + # Hidden dim for action model (auto-aligned at runtime) + "action_hidden_dim": 1024, + "hidden_size": 1024, + # Whether to add positional embeddings in the action head + "add_pos_embed": True, + "max_seq_len": 1024, + # Dimensionality of each action vector (e.g., 7 for 6-DoF + gripper) + "action_dim": 7, + # State dimension (proprioception input) + "state_dim": 7, + # Canonical chunk length (number of action steps the head predicts). + # Legacy YAMLs may use future_action_window_size = action_horizon - 1; + # apply_config_compat normalises both directions. + "action_horizon": 8, + # Repeat factor for flow-matching loss (more noise samples per batch) + "repeated_diffusion_steps": 8, + # Beta distribution params for noise schedule + "noise_beta_alpha": 1.5, + "noise_beta_beta": 1.0, + "noise_s": 0.999, + "num_timestep_buckets": 1000, + # Inference denoising steps + "num_inference_timesteps": 4, + # Number of vision tokens fed to action head + "num_target_vision_tokens": 32, + # === DiT Transformer sub-config === + "diffusion_model_cfg": { + # Cross-attention dim (aligned to VLM hidden_size at runtime) + "cross_attention_dim": 2048, + "dropout": 0.2, + "final_dropout": True, + "interleave_self_attention": True, + "norm_type": "ada_norm", + "num_layers": 16, + "output_dim": 1024, + "positional_embeddings": None, + }, + } + ) + + # # === Training precision flag === This is unnecessary, unused parameter + # reduce_in_full_precision: bool = True + + +@FRAMEWORK_REGISTRY.register("QwenGR00T") +class Qwen_GR00T(baseframework): + """ + Multimodal vision-language-action model (GR00T variant). + + Components: + - Qwen2.5-VL / Qwen3-VL backbone for fused language/vision token embeddings + - Flow-matching (DiT) diffusion head for continuous action sequence modeling + + Focus: Predict future continuous actions conditioned on images + instruction. + """ + + def __init__( + self, + config: Optional[dict] = None, + **kwargs, + ) -> None: + """ + Construct all submodules and cache key configuration values. + + Args: + config: Hierarchical configuration (OmegaConf/dict) containing framework + trainer sections. + **kwargs: Reserved for future overrides (unused). + """ + super().__init__() + # Merge framework defaults with YAML config (YAML wins on conflicts) + self.config = merge_framework_config(QwenGR00TDefaultConfig, config) + self.qwen_vl_interface = get_vlm_model(config=self.config) + # Florence conditions the DiT through its 1024-wide projection. + self.config.framework.action_model.diffusion_model_cfg.cross_attention_dim = ( + self.qwen_vl_interface.model.config.hidden_size + ) + + self.action_model: FlowmatchingActionHead = get_action_model(config=self.config) + # ZeRO-3 partitions all trainable parameters together, so the freshly + # initialized action head must match the BF16 Qwen backbone. + qwen_dtype = next(self.qwen_vl_interface.model.parameters()).dtype + self.action_model.to(dtype=qwen_dtype) + + # `action_horizon` is the single source of truth for chunk length. + # Legacy aliases (`future_action_window_size`, `past_action_window_size`) + # are normalised upstream by `share_tools.apply_config_compat`, so we + # only ever read `action_horizon` here. + self.action_horizon = int(self.config.framework.action_model.action_horizon) + + def forward( + self, + examples: List[dict] = None, + **kwargs, + ) -> Tuple: + """ """ + return_per_sample_loss = bool(kwargs.pop("return_per_sample_loss", False)) + flow_noise_seed = kwargs.pop("flow_noise_seed", None) + return_flow_noise_time = bool(kwargs.pop("return_flow_noise_time", False)) + if kwargs: + raise TypeError(f"Unexpected forward arguments: {sorted(kwargs)}") + batch_images = [example["image"] for example in examples] # [B，[PLT]] + instructions = [example["lang"] for example in examples] # [B, str] + actions = [example["action"] for example in examples] # label [B， len, 7] + + state = [example["state"] for example in examples] if "state" in examples[0] else None # [B, 1, state_dim] + + # Step 1: QWenVL input format + qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions) + with torch.autocast("cuda", dtype=torch.bfloat16): + qwenvl_outputs = self.qwen_vl_interface( + **qwen_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + # last_hidden_state: [B, seq_len, H] + last_hidden = qwenvl_outputs.hidden_states[-1] # [B, L, H] + + # Step 4: Action Expert Forward and Loss + with torch.autocast("cuda", dtype=torch.float32): + actions = torch.tensor( + np.array(actions), device=last_hidden.device, dtype=last_hidden.dtype + ) # [B, T_full, action_dim] + actions_target = actions[:, -self.action_horizon :, :] # (B, action_horizon, action_dim) + + repeated_diffusion_steps = int( + self.config.trainer.repeated_diffusion_steps + ) + actions_target_repeated = actions_target.repeat(repeated_diffusion_steps, 1, 1) + last_hidden_repeated = last_hidden.repeat(repeated_diffusion_steps, 1, 1) + state_repeated = None + if state is not None: + state = torch.tensor(np.array(state), device=last_hidden.device, dtype=last_hidden.dtype) + state_repeated = state.repeat(repeated_diffusion_steps, 1, 1) + + explicit_flow = [ + "loss_is_flow_noise" in example and "loss_is_flow_time" in example + for example in examples + ] + if any(explicit_flow) and not all(explicit_flow): + raise ValueError("A batch cannot mix explicit and sampled flow realizations") + flow_noise = None + flow_time = None + if all(explicit_flow): + flow_noise = torch.stack( + [ + torch.as_tensor(example["loss_is_flow_noise"]) + for example in examples + ], + dim=1, + ).reshape(actions_target_repeated.shape) + flow_time = torch.stack( + [ + torch.as_tensor(example["loss_is_flow_time"]) + for example in examples + ], + dim=1, + ).reshape(actions_target_repeated.shape[0]) + + action_loss = self.action_model( + last_hidden_repeated, actions_target_repeated, state_repeated, + reduction="none" if return_per_sample_loss else "mean", + flow_noise_seed=flow_noise_seed, + flow_noise=flow_noise, + flow_time=flow_time, + return_flow_noise_time=return_flow_noise_time, + ) # (B, chunk_len, action_dim) + sampled_noise = None + sampled_time = None + if return_flow_noise_time: + action_loss, sampled_noise, sampled_time = action_loss + if return_per_sample_loss: + action_loss = aggregate_repeated_losses( + action_loss, + repeat_count=repeated_diffusion_steps, + batch_size=len(examples), + ) + + output = {"action_loss": action_loss} + if return_flow_noise_time: + output["flow_noise"] = sampled_noise + output["flow_time"] = sampled_time + return output + + @torch.inference_mode() + def predict_action( + self, + examples: List[dict], + **kwargs: str, + ) -> np.ndarray: + """ + Steps: + 1. Resize images to training resolution (if specified) + 2. Encode with QwenVL (hidden states retained) + 6. Return normalized action trajectory + Returns: + dict: + normalized_actions (np.ndarray): Shape [B, T, action_dim], diffusion-sampled normalized actions. + """ + if type(examples) is not list: + examples = [examples] + batch_images = [to_pil_preserve(example["image"]) for example in examples] # [B，[PLT]] + instructions = [example["lang"] for example in examples] # [B, str] + + state = [example["state"] for example in examples] if "state" in examples[0] else None # [B, 1, state_dim] + + train_obs_image_size = getattr(self.config.datasets.vla_data, "image_size", None) + if train_obs_image_size: + batch_images = resize_images(batch_images, target_size=train_obs_image_size) + + # Step 1: QWenVL input format + qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions) + with torch.autocast("cuda", dtype=torch.bfloat16): + qwenvl_outputs = self.qwen_vl_interface( + **qwen_inputs, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + + # last_hidden_state: [B, seq_len, H] + last_hidden = qwenvl_outputs.hidden_states[-1] # [B, L, H] + + state = ( + torch.from_numpy(np.array(state)).to(last_hidden.device, dtype=last_hidden.dtype) + if state is not None + else None + ) + + # Step 4: Action Expert Forward + with torch.autocast("cuda", dtype=torch.float32): + pred_actions = self.action_model.predict_action(last_hidden, state) + + normalized_actions = pred_actions.detach().cpu().numpy() + return {"normalized_actions": normalized_actions} + + +if __name__ == "__main__": + import argparse + import os + + from omegaconf import OmegaConf + + parser = argparse.ArgumentParser() + parser.add_argument( + "--config_yaml", + type=str, + default="examples/LIBERO/train_files/starvla_cotrain_libero.yaml", + help="Path to YAML config", + ) + args, clipargs = parser.parse_known_args() + + if os.getenv("DEBUGPY_ENABLE", "0") == "1": + import debugpy + + debugpy.listen(("0.0.0.0", 10092)) + print("Rank 0 waiting for debugger attach on port 10092...") + debugpy.wait_for_client() + + cfg = OmegaConf.load(args.config_yaml) + + model: Qwen_GR00T = Qwen_GR00T(cfg) + print(model) + + image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)) + sample = { + "action": np.random.uniform(-1, 1, size=(16, 7)).astype(np.float16), + "image": [image], + "lang": "This is a fake instruction for testing.", + } + sample2 = sample.copy() + sample2["lang"] = "Another fake instruction for testing." + + batch = [sample, sample2] + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device) + forward_output = model(batch) + action_loss = forward_output["action_loss"] + print(f"Action Loss: {action_loss.item()}") + + predict_output = model.predict_action(examples=[sample]) + normalized_actions = predict_output["normalized_actions"] + print(f"Unnormalized Action: {normalized_actions}") + + print("Finished") diff --git a/real_robot/source/starVLA/model/framework/VLM4A/__init__.py b/real_robot/source/starVLA/model/framework/VLM4A/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..81330e2ca135c0a5d25e94d8060a267b50661678 --- /dev/null +++ b/real_robot/source/starVLA/model/framework/VLM4A/__init__.py @@ -0,0 +1,8 @@ +"""VLM4A — VLM-for-Action frameworks. + +All VLA frameworks that use a Vision-Language Model (Qwen-VL, Florence, etc.) +as the perception backbone live here. + +The registry decorator on each class registers them in FRAMEWORK_REGISTRY, +so `build_framework(cfg)` can find them after auto-import. +""" diff --git a/real_robot/source/starVLA/model/framework/__init__.py b/real_robot/source/starVLA/model/framework/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cca1db73ef45e5732d73d74dc50a7da42d34096f --- /dev/null +++ b/real_robot/source/starVLA/model/framework/__init__.py @@ -0,0 +1,8 @@ +"""Framework package internals. + +Public framework APIs are defined in `starVLA.model.framework.base_framework`. +""" + +from starVLA.model.tools import FRAMEWORK_REGISTRY + +__all__ = ["FRAMEWORK_REGISTRY"] diff --git a/real_robot/source/starVLA/model/framework/base_framework.py b/real_robot/source/starVLA/model/framework/base_framework.py new file mode 100644 index 0000000000000000000000000000000000000000..47497c0bea94f3dd971d670967d5d80a6ffd7acf --- /dev/null +++ b/real_robot/source/starVLA/model/framework/base_framework.py @@ -0,0 +1,290 @@ +# Modified in the ReTime research artifact; see NOTICE. +""" +Base framework abstraction providing: +- Pretrained loading (config + normalization stats + weights) +- Action space utilities (dimension, stats, (un)normalization) +- Trainable module discovery helper +Note: No device placement or optimizer concerns handled here (delegated to trainer). +""" + +import importlib +import pkgutil +from pathlib import Path +from typing import Any, Dict, List + +import numpy as np +import torch +from transformers import PretrainedConfig, PreTrainedModel + +from starVLA.model.framework.share_tools import dict_to_namespace, read_mode_config +from starVLA.model.tools import FRAMEWORK_REGISTRY, FrameworkTools, auto_get_trainable_modules +from starVLA.training.trainer_utils import initialize_overwatch + +logger = initialize_overwatch(__name__) +_FRAMEWORKS_IMPORTED = False +_FRAMEWORK_IMPORT_ERRORS: dict[str, str] = {} +_OPTIONAL_EXTERNAL_MODULES = {"qwen_vl_utils"} + + +def _import_optional_framework(module_path: str) -> None: + """Import a framework plugin without making unrelated extras mandatory.""" + try: + importlib.import_module(module_path) + except ModuleNotFoundError as exc: + missing_root = (exc.name or "").split(".", 1)[0] + if missing_root not in _OPTIONAL_EXTERNAL_MODULES: + raise + _FRAMEWORK_IMPORT_ERRORS[module_path] = str(exc) + logger.warning("Skipping optional framework %s: %s", module_path, exc) + + +def _auto_import_framework_modules() -> None: + global _FRAMEWORKS_IMPORTED + if _FRAMEWORKS_IMPORTED: + return + + _SKIP = {"__init__", "base_framework", "share_tools"} + framework_dir = Path(__file__).resolve().parent + + # Scan top-level modules (backwards compat) + for _, module_name, is_pkg in pkgutil.iter_modules([str(framework_dir)]): + if module_name in _SKIP: + continue + if is_pkg: + # Scan sub-packages (VLM4A/, WM4A/, etc.) + sub_dir = framework_dir / module_name + for _, sub_name, _ in pkgutil.iter_modules([str(sub_dir)]): + if sub_name.startswith("_"): + continue + _import_optional_framework( + f"starVLA.model.framework.{module_name}.{sub_name}" + ) + else: + _import_optional_framework(f"starVLA.model.framework.{module_name}") + + _FRAMEWORKS_IMPORTED = True + + +def build_framework(cfg): # The single entry point for building different model frameworks + """ + Build a framework model from config. + Args: + cfg: Config object containing `cfg.framework.name`. + Returns: + nn.Module: Instantiated framework model. + """ + if not hasattr(cfg, "framework") or not hasattr(cfg.framework, "name"): + raise ValueError("Missing `cfg.framework.name`. The framework API now only accepts `framework.name`.") + + _auto_import_framework_modules() + + framework_id = cfg.framework.name + if framework_id not in FRAMEWORK_REGISTRY._registry: + available = sorted(FRAMEWORK_REGISTRY._registry.keys()) + skipped = "; ".join( + f"{module}: {error}" + for module, error in sorted(_FRAMEWORK_IMPORT_ERRORS.items()) + ) + raise NotImplementedError( + f"Framework `{framework_id}` is not implemented. Available frameworks: {available}. " + f"Skipped optional frameworks: {skipped or 'none'}" + ) + + model_class = FRAMEWORK_REGISTRY[framework_id] + return model_class(cfg) + + +# PreTrainedModel, AutoModel, PretrainedConfig, are so good, find sometime to study them +# Framework construction merges the YAML and transformer configurations. + + +class baseframework(PreTrainedModel): + """ + Lightweight base class for higher-level VLA model assemblies. + Subclasses are expected to: + - Accept a structured config + - Register components in __init__ + - Use provided helpers for action normalization handling + """ + + def __init__(self, hf_config=PretrainedConfig()) -> None: + """ + Initialize base nn.Module. Subclasses add components. + """ + + super().__init__(hf_config) + + # ------------------------------------------------------------------ + # Soft-constraint interface: subclasses should override these. + # Default implementations raise NotImplementedError so that IDE + # tooling (e.g. pylance, mypy) flags missing overrides, while + # still allowing PreTrainedModel instantiation (no ABC). + # ------------------------------------------------------------------ + + def forward(self, examples: List[dict], **kwargs) -> dict: + """Training forward pass. + + Args: + examples: List[dict], each dict requires at least: + - image: List[PIL.Image] + - lang: str + - action: np.ndarray shaped [T, action_dim] + + Returns: + dict: Must contain ``"action_loss"`` (torch.Tensor scalar). + May contain extra keys for logging (e.g. ``"kl_loss"``). + """ + raise NotImplementedError( + f"{type(self).__name__} must implement forward(examples) -> dict with 'action_loss' key." + ) + + def predict_action(self, examples: List[dict], **kwargs) -> dict: + """Inference: predict future actions from observations. + + Args: + examples: Same schema as *forward* (minus ``action`` which is optional). + **kwargs: Framework-specific inference options (e.g. ``use_ddim``). + + Returns: + dict: Must contain ``"normalized_actions"`` (np.ndarray [B, T, action_dim]). + """ + raise NotImplementedError( + f"{type(self).__name__} must implement predict_action(examples) -> dict with 'normalized_actions' key." + ) + + # ------------------------------------------------------------------ + # Unified loss interface for Trainer + # ------------------------------------------------------------------ + + def supports_training_tag(self, tag: str) -> bool: + """Return whether this framework can consume batches for *tag*.""" + if tag == "vla": + return type(self).forward is not baseframework.forward + if tag == "vlm": + return hasattr(self, "qwen_vl_interface") or type(self).forward_vlm is not baseframework.forward_vlm + return False + + def compute_loss(self, tag: str, batch, loss_scale: dict = None) -> Dict[str, torch.Tensor] | None: + """Unified forward entry-point: route to the right forward by *tag*. + + The trainer calls ``model.compute_loss(tag, batch)`` for every + ``(tag, batch)`` pair produced by :class:`DataLoaderManager`. + The model internally dispatches: + + - ``"vla"`` → ``self.forward(batch)`` + - ``"vlm"`` → ``self.forward_vlm(batch)`` + + Subclasses can override this to add more tags (e.g. ``"world"``). + + Args: + tag: dataset type tag (``"vla"``, ``"vlm"``, …) + batch: the batch produced by the corresponding DataLoader. + loss_scale: ``{"vla": 1.0, "vlm": 0.1}`` per-tag loss multiplier. + Defaults to 1.0 for unspecified tags. + + Returns: + dict[str, Tensor] | None: keyed losses (e.g. ``{"action_loss": ...}``). + Returns ``None`` when this framework does not support the + incoming dataloader tag so the trainer can ``continue``. + """ + if not self.supports_training_tag(tag): + return None + + scale = (loss_scale or {}).get(tag, 1.0) + + if tag == "vla": + out = self.forward(batch) + elif tag == "vlm": + out = self.forward_vlm(batch) + else: + return None + + # Apply loss scale and filter to Tensor values only + return {k: v * scale for k, v in out.items() if isinstance(v, torch.Tensor)} + + def forward_vlm(self, batch) -> Dict[str, torch.Tensor]: + """VLM forward pass (default implementation). + + Delegates to ``self.qwen_vl_interface(**batch)`` which is present on + every framework subclass that uses a Qwen VL backbone. + + Subclasses may override to add custom VLM logic. + + Args: + batch: dict produced by the VLM dataloader. + + Returns: + dict: Must contain ``"vlm_loss"`` (torch.Tensor scalar). + """ + if not hasattr(self, "qwen_vl_interface"): + raise NotImplementedError( + f"{type(self).__name__} has no `qwen_vl_interface`. " + "Override forward_vlm() to support VLM training." + ) + out = self.qwen_vl_interface(**batch) + return {"vlm_loss": out.loss} + + @classmethod + def from_pretrained( + cls, + pretrained_checkpoint: str, + **kwargs, + ) -> None: + """ + Restore a model instance from a saved checkpoint. + + Workflow: + 1. Resolve checkpoint path + 2. Load config + dataset normalization statistics + 3. Build model with loaded config + 4. Load state_dict strictly (reports missing/unexpected keys) + 5. Attach normalization stats for later un-normalization + + Args: + pretrained_checkpoint: Path to .pt file inside run/checkpoints directory. + **kwargs: Extra constructor overrides passed to subclass. + + Returns: + baseframework: Instantiated model (left on CPU; caller decides device). + + Raises: + RuntimeError: If state_dict key mismatch occurs under strict=True. + FileNotFoundError: If underlying files are missing (surfaced earlier). + """ + pretrained_checkpoint = Path(pretrained_checkpoint) + model_config, norm_stats = read_mode_config(pretrained_checkpoint) # read config and norm_stats + + config = dict_to_namespace(model_config) + model_config = config + model_config.trainer.pretrained_checkpoint = None + + FrameworkModel = build_framework(cfg=model_config) + # set for action un-norm + FrameworkModel.norm_stats = norm_stats + # Load from Checkpoint (Custom --> should load both *projector* and *llm* weights) + if pretrained_checkpoint.suffix == ".safetensors": + from safetensors.torch import load_file + + model_state_dict = load_file(str(pretrained_checkpoint)) + else: + model_state_dict = torch.load(pretrained_checkpoint, map_location="cpu") + # logger.info(f"Loading model weights from `{pretrained_checkpoint}`") + model_keys = set(FrameworkModel.state_dict().keys()) + checkpoint_keys = set(model_state_dict.keys()) + try: + FrameworkModel.load_state_dict(model_state_dict, strict=True) + except RuntimeError as e: + # must keep all keys matched + common_keys = model_keys.intersection(checkpoint_keys) + missing_keys = model_keys - common_keys + unexpected_keys = checkpoint_keys - common_keys + if missing_keys: + logger.warning(f"Missing keys in state_dict: {missing_keys}") + if unexpected_keys: + logger.warning(f"Unexpected keys in state_dict: {unexpected_keys}") + + raise e + + # **ensure model is on GPU** + FrameworkModel = FrameworkModel + return FrameworkModel diff --git a/real_robot/source/starVLA/model/framework/share_tools.py b/real_robot/source/starVLA/model/framework/share_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..a8c8c4291065acd84096dbf1703292f37dc73a70 --- /dev/null +++ b/real_robot/source/starVLA/model/framework/share_tools.py @@ -0,0 +1,550 @@ +# Modified in the ReTime research artifact; see NOTICE. +""" +Shared configuration / utility helpers for framework components: +- NamespaceWithGet: lightweight namespace behaving like a dict +- OmegaConf conversion helpers +- Config merging decorator for model __init__ +- Checkpoint config/statistics loading +""" + +import functools +import inspect +import json +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from omegaconf import OmegaConf + +from starVLA.training.trainer_utils import initialize_overwatch + +# Initialize Overwatch =>> Wraps `logging.Logger` +overwatch = initialize_overwatch(__name__) + + +class NamespaceWithGet(SimpleNamespace): + def get(self, key, default=None): + """ + Return attribute value if present, else default (dict-like API). + + Args: + key: Attribute name. + default: Fallback if attribute missing. + + Returns: + Any: Stored value or default. + """ + return getattr(self, key, default) + + def items(self): + """ + Iterate (key, value) pairs like dict.items(). + + Returns: + Generator[Tuple[str, Any], None, None] + """ + return ((key, getattr(self, key)) for key in self.__dict__) + + def __iter__(self): + """ + Return iterator over attribute keys (enables dict unpacking **obj). + + Returns: + Iterator[str] + """ + return iter(self.__dict__) + + def to_dict(self): + """ + Recursively convert nested NamespaceWithGet objects into plain dicts. + + Returns: + dict: Fully materialized dictionary structure. + """ + return {key: value.to_dict() if isinstance(value, NamespaceWithGet) else value for key, value in self.items()} + + +def dict_to_namespace(d): + """ + Create an OmegaConf config from a plain dictionary. + + Args: + d: Input dictionary. + + Returns: + OmegaConf: DictConfig instance. + """ + return OmegaConf.create(d) + + +def _to_omegaconf(x: Any): + """ + Convert diverse input types into an OmegaConf object. + + Accepted types: + - None -> empty DictConfig + - str path -> load YAML/JSON via OmegaConf.load + - dict -> DictConfig + - DictConfig / ListConfig -> returned unchanged + - NamespaceWithGet / SimpleNamespace -> converted via vars()/to_dict() + + Args: + x: Input candidate. + + Returns: + OmegaConf: Normalized configuration node. + """ + if x is None: + return OmegaConf.create({}) + if isinstance(x, OmegaConf.__class__): # fallback, typically not hit + return x + try: + # OmegaConf node detection + from omegaconf import DictConfig, ListConfig + + if isinstance(x, (DictConfig, ListConfig)): + return x + except Exception: + pass + + if isinstance(x, str): + # treat as path + return OmegaConf.load(x) + if isinstance(x, dict): + return OmegaConf.create(x) + if isinstance(x, NamespaceWithGet) or isinstance(x, SimpleNamespace): + # convert to plain dict + try: + d = x.to_dict() if hasattr(x, "to_dict") else vars(x) + except Exception: + d = vars(x) + return OmegaConf.create(d) + # fallback: try to create + return OmegaConf.create(x) + + +def merge_pram_config(init): + """ + Decorator for __init__ to unify config handling. + + Behavior: + 1. Extract 'config' kwarg / arg (path | dict | OmegaConf | namespace) + 2. Convert to OmegaConf + 3. Merge with explicitly passed init parameters (explicit overrides file) + 4. Attach merged config to self.config + 5. Call original __init__ with merged config + + Args: + init: Original __init__ function. + + Returns: + Wrapped initializer. + """ + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + # Map positional args to parameter names (excluding self) + sig = inspect.signature(init) + param_names = [name for i, (name, p) in enumerate(sig.parameters.items()) if i > 0] + + init_kwargs = {} + for name, val in zip(param_names, args): + init_kwargs[name] = val + # override with explicit kwargs + init_kwargs.update(kwargs) + + # get provided config (if any) + provided_config = init_kwargs.get("config", None) + + loaded_cfg = _to_omegaconf(provided_config) + + # build params cfg from explicit init args (other than config) + params = {k: v for k, v in init_kwargs.items() if k != "config"} + params_cfg = OmegaConf.create(params) if params else OmegaConf.create({}) + + # merge: loaded_cfg <- params_cfg (params override file) + merged = OmegaConf.merge(loaded_cfg, params_cfg) + + # set on instance + try: + # prefer attaching OmegaConf directly + self.config = merged + except Exception: + # fallback to dict + self.config = OmegaConf.to_container(merged, resolve=True) + + # prepare kwargs for original init: ensure config is the merged OmegaConf + call_kwargs = dict(init_kwargs) + call_kwargs["config"] = merged + + # call original __init__ using keyword args only (safer) + return init(self, **call_kwargs) + + return wrapper + + +def merge_framework_config(default_config_cls, cfg): + """ + Merge a framework's default config (dataclass) with the incoming YAML config. + + Rules: + - default_config_cls provides documented defaults for `cfg.framework` + - YAML values (cfg.framework) override matching defaults + - Extra YAML keys not in defaults are preserved (Config-as-API flexibility) + - Missing YAML keys fall back to defaults (less YAML boilerplate) + + The merge only touches the `cfg.framework` sub-tree; datasets / trainer / etc. + are left untouched. + + Args: + default_config_cls: A dataclass **class** (not instance) whose fields() define + the default framework config with type hints and comments. + cfg: The full OmegaConf config (must contain cfg.framework). + + Returns: + cfg: The same config object with cfg.framework replaced by the merged result. + """ + import dataclasses + + from omegaconf import DictConfig, OmegaConf + + # 1. Instantiate defaults and convert to OmegaConf + defaults_instance = default_config_cls() + defaults_dict = dataclasses.asdict(defaults_instance) + defaults_omega = OmegaConf.create(defaults_dict) + + # 2. Extract the YAML framework section + if hasattr(cfg, "framework"): + # Unwrap AccessTrackedConfig if needed + yaml_fw = cfg.framework + if hasattr(yaml_fw, "_cfg"): + yaml_fw = yaml_fw._cfg + if not isinstance(yaml_fw, DictConfig): + yaml_fw = OmegaConf.create(yaml_fw if isinstance(yaml_fw, dict) else {}) + else: + yaml_fw = OmegaConf.create({}) + + # 3. Merge: defaults first, YAML overrides (YAML wins on conflicts) + merged_fw = OmegaConf.merge(defaults_omega, yaml_fw) + + # 4. Write back into the original cfg + # Handle both OmegaConf and AccessTrackedConfig transparently + if hasattr(cfg, "_cfg") and isinstance(cfg._cfg, DictConfig): + # AccessTrackedConfig caches child wrappers in _children dict. + # After replacing the underlying DictConfig node, the old child wrapper + # still points to the pre-merge node (stale data). We must invalidate + # the cache so the next attribute access creates a fresh wrapper around + # the merged node. + # + # However, the old child's _local_accessed set records which keys were + # already read (e.g. "name" from build_framework). Deleting the child + # would lose that tracking info, causing save_accessed_config to omit + # those keys from config.yaml. So we preserve and restore it. + cfg._cfg.framework = merged_fw + if hasattr(cfg, "_children") and "framework" in cfg._children: + old_accessed = cfg._children["framework"]._local_accessed.copy() + del cfg._children["framework"] # invalidate stale cache + new_child = cfg.framework # re-create child around merged_fw + new_child._local_accessed.update(old_accessed) # restore tracking + elif isinstance(cfg, DictConfig): + cfg.framework = merged_fw + else: + # Fallback — try direct attribute setting + try: + cfg.framework = merged_fw + except Exception: + overwatch.warning("Could not write merged framework config back to cfg.") + + return cfg + + +def populate_layerwise_dit_cfg(cfg, *, dit_hidden_dim: int, num_dit_layers: int): + """ + Populate ``framework.action_model.diffusion_model_cfg`` with the DiT shape + fields required by ``LayerwiseFlowmatchingActionHead``. + + Why this helper exists: + The action head is intentionally agnostic of the VLM backbone — it only + consumes ``diffusion_model_cfg``. Each framework (QwenPI, QwenPI_v3, + ...) is responsible for deciding the DiT shape (depth + hidden) from + whatever source it likes (LLM hidden, a compressed projector dim, ...) + and writing it here BEFORE calling ``get_action_model``. + + Fields written (override any stale YAML values): + - num_layers = num_dit_layers + - input_embedding_dim = dit_hidden_dim + - cross_attention_dim = dit_hidden_dim (encoder is pre-projected) + - num_attention_heads = dit_hidden_dim // attention_head_dim + (uses existing attention_head_dim if set, else 64) + + Args: + cfg: Full OmegaConf config. + dit_hidden_dim: DiT internal hidden dim. + num_dit_layers: Number of DiT cross-attention layers. + + Returns: + The (mutated) diffusion_model_cfg node. + """ + dit_cfg = cfg.framework.action_model.diffusion_model_cfg + head_dim = dit_cfg.get("attention_head_dim", None) or 64 + dit_cfg.attention_head_dim = head_dim + dit_cfg.num_layers = int(num_dit_layers) + dit_cfg.input_embedding_dim = int(dit_hidden_dim) + dit_cfg.cross_attention_dim = int(dit_hidden_dim) + dit_cfg.num_attention_heads = int(dit_hidden_dim) // int(head_dim) + return dit_cfg + + +def read_model_config(pretrained_checkpoint): + """ + Load global model configuration and dataset normalization statistics + associated with a saved checkpoint (.pt). + + Expected directory layout: + /checkpoints/.pt + /config.json + /dataset_statistics.json + + Args: + pretrained_checkpoint: Path to a .pt checkpoint file. + + Returns: + tuple: + global_cfg (dict): Loaded config.json contents. + norm_stats (dict): Dataset statistics for (de)normalization. + + Raises: + FileNotFoundError: If checkpoint or required JSON files are missing. + AssertionError: If file suffix or structure invalid. + """ + if os.path.isfile(pretrained_checkpoint): + overwatch.info(f"Loading from local checkpoint path `{(checkpoint_pt := Path(pretrained_checkpoint))}`") + + # [Validate] Checkpoint Path should look like + # `...//checkpoints/.pt|.safetensors` + assert checkpoint_pt.suffix in {".pt", ".safetensors"} + run_dir = checkpoint_pt.parents[1] + + # Get paths for `config.json`, `dataset_statistics.json` and pretrained checkpoint + config_json, dataset_statistics_json = run_dir / "config.json", run_dir / "dataset_statistics.json" + assert config_json.exists(), f"Missing `config.json` for `{run_dir = }`" + assert dataset_statistics_json.exists(), f"Missing `dataset_statistics.json` for `{run_dir = }`" + + # Otherwise =>> try looking for a match on `model_id_or_path` on the HF Hub (`model_id_or_path`) + # Load VLA Config (and corresponding base VLM `ModelConfig`) from `config.json` + with open(config_json, "r") as f: + global_cfg = json.load(f) + + # Normalise legacy / pre-v0.21 configs to current schema (idempotent; + # also ensures `past_action_window_size`, `action_horizon`, + # `future_action_window_size`, etc. are all present for downstream code). + try: + _oc = OmegaConf.create(global_cfg) + apply_config_compat(_oc) + global_cfg = OmegaConf.to_container(_oc, resolve=True) + except Exception as e: + overwatch.warning(f"apply_config_compat failed on `{config_json}`: {e}") + + # Load Dataset Statistics for Action Denormalization + with open(dataset_statistics_json, "r") as f: + norm_stats = json.load(f) + else: + overwatch.error(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + raise FileNotFoundError(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + return global_cfg, norm_stats + + +def read_mode_config(pretrained_checkpoint): + """ + Same as read_model_config (legacy duplicate kept for backward compatibility). + + Args: + pretrained_checkpoint: Path to a .pt checkpoint file. + + Returns: + tuple: + vla_cfg (dict) + norm_stats (dict) + """ + if os.path.isfile(pretrained_checkpoint): + overwatch.info(f"Loading from local checkpoint path `{(checkpoint_pt := Path(pretrained_checkpoint))}`") + + # [Validate] Checkpoint Path should look like + # `...//checkpoints/.pt|.safetensors` + assert checkpoint_pt.suffix in {".pt", ".safetensors"} + run_dir = checkpoint_pt.parents[1] + + # Get paths for `config.json`, `dataset_statistics.json` and pretrained checkpoint + config_yaml, dataset_statistics_json = run_dir / "config.yaml", run_dir / "dataset_statistics.json" + assert config_yaml.exists(), f"Missing `config.yaml` for `{run_dir = }`" + assert dataset_statistics_json.exists(), f"Missing `dataset_statistics.json` for `{run_dir = }`" + + # Otherwise =>> try looking for a match on `model_id_or_path` on the HF Hub (`model_id_or_path`) + # Load VLA Config (and corresponding base VLM `ModelConfig`) from `config.json` + try: + ocfg = OmegaConf.load(str(config_yaml)) + # Normalise legacy / pre-v0.21 configs to current schema (idempotent). + apply_config_compat(ocfg) + global_cfg = OmegaConf.to_container(ocfg, resolve=True) + except Exception as e: + overwatch.error(f"Failed to load YAML config `{config_yaml}`: {e}") + raise + + # Load Dataset Statistics for Action Denormalization + with open(dataset_statistics_json, "r") as f: + norm_stats = json.load(f) + else: + overwatch.error(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + raise FileNotFoundError(f"Pretrained checkpoint `{pretrained_checkpoint}` does not exist.") + return global_cfg, norm_stats + + +# Normalize legacy configuration keys before model construction. The +# transformation is idempotent and leaves framework class signatures intact. + +CONFIG_VERSION = "0.21" + + +def apply_config_compat(cfg, *, strict: bool = False): + """ + Normalise an arbitrary (old or new) starVLA training config into the + current `version_id == "0.21"` schema. + + Performed transformations (each applied only when needed): + + 1. `framework.action_model.action_horizon` ↔ `future_action_window_size` + - `action_horizon` is canonical (preferred user-facing name). + - `future_action_window_size = action_horizon - 1` is auto-filled so + framework code that still reads the old key keeps working. + - If both are present and inconsistent, a warning is emitted and + `action_horizon` wins. + + 2. `framework.action_model.diffusion_model_cfg.output_dim` + - Auto-filled from `framework.action_model.hidden_size` when missing. + + 3. `framework.action_model.diffusion_model_cfg.cross_attention_dim` + - Auto-filled from `framework.qwenvl.vl_hidden_dim` when missing. + Frameworks that further override this at runtime (e.g. QwenGR00T) + are unaffected. + + 4. `framework.action_model.action_hidden_dim` + - Auto-filled from `hidden_size` when missing. OFT-family frameworks + still overwrite this from VLM hidden_size at runtime. + + 5. `framework.action_model.past_action_window_size` + - Auto-filled to `0` when missing. All released starVLA frameworks + run with past=0; the field is therefore dropped from user YAMLs + and only re-materialised here for legacy code that still reads it. + + 6. `cfg.version_id` is stamped to `"0.21"`. + + Args: + cfg: An OmegaConf DictConfig (or anything _to_omegaconf can wrap). + strict: If True, raise on inconsistent values instead of warning. + + Returns: + The same `cfg` object (mutated in place) for chaining convenience. + """ + from omegaconf import OmegaConf + + if cfg is None: + return cfg + + src_version = OmegaConf.select(cfg, "version_id", default=None) + + # ---- 1. action_horizon ↔ future_action_window_size ---- + am_path = "framework.action_model" + am = OmegaConf.select(cfg, am_path, default=None) + if am is not None: + ah = OmegaConf.select(am, "action_horizon", default=None) + fw = OmegaConf.select(am, "future_action_window_size", default=None) + + if ah is None and fw is not None: + ah = int(fw) + 1 + OmegaConf.update(cfg, f"{am_path}.action_horizon", ah, force_add=True) + elif ah is not None and fw is None: + fw = int(ah) - 1 + OmegaConf.update(cfg, f"{am_path}.future_action_window_size", fw, force_add=True) + elif ah is not None and fw is not None and int(ah) != int(fw) + 1: + msg = ( + f"[apply_config_compat] inconsistent action_horizon={ah} vs " + f"future_action_window_size={fw}; expected action_horizon == future + 1. " + "Using action_horizon as canonical." + ) + if strict: + raise ValueError(msg) + overwatch.warning(msg) + OmegaConf.update(cfg, f"{am_path}.future_action_window_size", int(ah) - 1, force_add=True) + + # ---- 2 & 3. diffusion_model_cfg auto-fill ---- + dm_path = f"{am_path}.diffusion_model_cfg" + dm = OmegaConf.select(cfg, dm_path, default=None) + if dm is not None: + hidden_size = OmegaConf.select(am, "hidden_size", default=None) + if OmegaConf.select(dm, "output_dim", default=None) is None and hidden_size is not None: + OmegaConf.update(cfg, f"{dm_path}.output_dim", int(hidden_size), force_add=True) + + if OmegaConf.select(dm, "cross_attention_dim", default=None) is None: + vl_hidden = OmegaConf.select(cfg, "framework.qwenvl.vl_hidden_dim", default=None) + if vl_hidden is not None: + OmegaConf.update(cfg, f"{dm_path}.cross_attention_dim", int(vl_hidden), force_add=True) + # else: leave None — framework __init__ may auto-bind it + + # ---- 4. action_hidden_dim fallback ---- + if OmegaConf.select(am, "action_hidden_dim", default=None) is None: + hidden_size = OmegaConf.select(am, "hidden_size", default=None) + if hidden_size is not None: + OmegaConf.update(cfg, f"{am_path}.action_hidden_dim", int(hidden_size), force_add=True) + + # ---- 5. past_action_window_size default ---- + if OmegaConf.select(am, "past_action_window_size", default=None) is None: + OmegaConf.update(cfg, f"{am_path}.past_action_window_size", 0, force_add=True) + + # ---- 6. stamp version ---- + if src_version != CONFIG_VERSION: + try: + OmegaConf.update(cfg, "version_id", CONFIG_VERSION, force_add=True) + except Exception: + try: + cfg.version_id = CONFIG_VERSION + except Exception: + pass + overwatch.info(f"[apply_config_compat] normalised config from version_id={src_version!r} to {CONFIG_VERSION!r}") + + return cfg + + +# ────────────────────────────────────────────────────────────────────── +# Discretised proprioceptive state → instruction prefix (π₀.5 style) +# ────────────────────────────────────────────────────────────────────── +import numpy as _np +from typing import List as _List + + +def state2str_transform(state: "_np.ndarray", num_bins: int = 256) -> str: + """Quantise a state vector into ``num_bins`` uniform bins over [-1, 1] + and return space-separated bin indices. + + Example: [-0.5, 0.1, 0.8] -> "95 133 203" + """ + discretized_state = _np.digitize(state, bins=_np.linspace(-1, 1, num_bins + 1)[:-1]) - 1 + return " ".join(map(str, discretized_state)) + + +def add_discretized_state_to_instruction( + instructions: "_List[str]", + states: "_List[_np.ndarray]", + num_bins: int = 256, +) -> "_List[str]": + """Append discretised proprioceptive state tokens to each instruction. + + Format: `` [STATE] [ACTION]`` + Lets the VLM attend to the robot state purely through its existing + text-token pathway — no extra encoder required (π₀.5 style). + """ + updated_instructions = [] + for instr, state in zip(instructions, states): + state_str = state2str_transform(state[0], num_bins=num_bins) + updated_instructions.append(f"{instr} [STATE] {state_str} [ACTION]") + return updated_instructions diff --git a/real_robot/source/starVLA/model/modules/action_model/GR00T_ActionHeader.py b/real_robot/source/starVLA/model/modules/action_model/GR00T_ActionHeader.py new file mode 100644 index 0000000000000000000000000000000000000000..9ca51428c8aa387a5ba8b85461a690c35dc8c832 --- /dev/null +++ b/real_robot/source/starVLA/model/modules/action_model/GR00T_ActionHeader.py @@ -0,0 +1,505 @@ +# Modified in the ReTime research artifact; see NOTICE. +# Copyright 2025 NVIDIA Corp. and affiliates. All rights reserved. +# Modification: [rm and add some connect adapter to match with starVLA, e.g., "rm "]. +# Action repeat is inspired by CogACT + + +from dataclasses import dataclass, field + +import torch +import torch.nn.functional as F +from torch import nn +from torch.distributions import Beta +from transformers import PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +from starVLA.model.modules.action_model.flow_matching_head.action_encoder import ( + SinusoidalPositionalEncoding, + swish, +) +from starVLA.model.modules.action_model.flow_matching_head.cross_attention_dit import DiT +from starVLA.training.infoflow import reduce_flow_matching_loss + +# TODO try to meger DiT Modules with follow_match_head, they are just the same arch, but diff loss, use diffusers package will be simple + + +class CategorySpecificLinear(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim): + super().__init__() + self.num_categories = num_categories + # For each category, we have separate weights and biases. + self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) + self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) + + def forward(self, x, cat_ids): + selected_W = self.W[cat_ids] + selected_b = self.b[cat_ids] + # import ipdb; ipdb.set_trace() + return torch.bmm(x, selected_W) + selected_b.unsqueeze(1) + + +class CategorySpecificMLP(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim, output_dim): + super().__init__() + self.num_categories = num_categories + self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) + self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) + + def forward(self, x, cat_ids): + hidden = F.relu(self.layer1(x, cat_ids)) + return self.layer2(hidden, cat_ids) + + +class MLP(nn.Module): + def __init__(self, input_dim, hidden_dim, output_dim): + super().__init__() + self.layer1 = nn.Linear(input_dim, hidden_dim) + self.layer2 = nn.Linear(hidden_dim, output_dim) + + def forward(self, x): + return self.layer2(F.relu(self.layer1(x))) + + +class ActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size): + super().__init__() + self.hidden_size = hidden_size + self.action_dim = action_dim + self.layer1 = nn.Linear(action_dim, hidden_size) + self.layer2 = nn.Linear(2 * hidden_size, hidden_size) + self.layer3 = nn.Linear(hidden_size, hidden_size) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.layer1(actions) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then layer2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.layer2(x)) + + # 5) Finally W3 => (B, T, w) + x = self.layer3(x) + return x + + +class MultiEmbodimentActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size, num_embodiments): + super().__init__() + self.hidden_size = hidden_size + self.num_embodiments = num_embodiments + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) + self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps, cat_ids): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + cat_ids: shape (B,) + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions, cat_ids) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x, cat_ids)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x, cat_ids) + return x + + +@dataclass +class FlowmatchingActionHeadConfig(PretrainedConfig): + """NOTE: N1.5 uses XEmbFlowmatchingPolicyHeadConfig as action head""" + + add_pos_embed: bool = field(default=True, metadata={"help": "Whether to add positional embedding"}) + diffusion_model_cfg: dict = field(default=None, metadata={"help": "Diffusion model configuration."}) + input_embedding_dim: int = field(default=1536, metadata={"help": "Input embedding channel dimension."}) + + hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."}) + max_seq_len: int = field(default=1024, metadata={"help": "Maxium Sequence Length"}) + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + noise_beta_alpha: float = field(default=1.5, metadata={"help": ""}) + noise_beta_beta: float = field(default=1.0, metadata={"help": ""}) + noise_s: float = field(default=0.999, metadata={"help": "Flow matching noise Beta distribution s."}) + num_timestep_buckets: int = field(default=1000, metadata={"help": "Number of timestep discretization buckets."}) + num_inference_timesteps: int = field( + default=None, + metadata={"help": "Number of inference steps for noise diffusion."}, + ) + max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."}) + tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."}) + tune_diffusion_model: bool = field(default=True, metadata={"help": "Whether to tune the diffusion model."}) + load_pretrained_det_decode_layer_path: str = field( + default=None, metadata={"help": "Path to pretrained detection model."} + ) + detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."}) + + freeze_decode_layer: bool = field(default=False) + expand_batch: int = field(default=None) + use_vlln: bool = field(default=True) + + vl_self_attention_cfg: dict = field(default=None) + num_target_vision_tokens: int = field(default=32, metadata={"help": "Number of target vision tokens."}) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +DiTConfig = { + "DiT-B": {"input_embedding_dim": 768, "attention_head_dim": 64, "num_attention_heads": 12}, + "DiT-L": {"input_embedding_dim": 1536, "attention_head_dim": 48, "num_attention_heads": 32}, +} + + +class FlowmatchingActionHead(nn.Module): + def __init__( + self, + full_config, + ): + super().__init__() + config = full_config.framework.action_model + self.full_config = full_config + + # ------------------------------------------------------------------ + # DiT architecture selection + # action_model_type: "DiT-B" | "DiT-L" + # DiT-B → input_embedding_dim=768, heads=12, head_dim=64 + # DiT-L → input_embedding_dim=1536, heads=32, head_dim=48 + # diffusion_model_cfg overrides/extends the base DiT shape. + # In particular, diffusion_model_cfg.cross_attention_dim MUST be + # set by the framework to match the VLM hidden size BEFORE calling + # get_action_model(), e.g.: + # cfg.framework.action_model.diffusion_model_cfg.cross_attention_dim + # = vlm.model.config.hidden_size + # ------------------------------------------------------------------ + action_model_type = config.action_model_type + action_model_cfg = DiTConfig[action_model_type] + self.input_embedding_dim = action_model_cfg["input_embedding_dim"] + + diffusion_model_cfg = config.diffusion_model_cfg + diffusion_model_cfg = {**action_model_cfg, **diffusion_model_cfg} + self.model = DiT(**diffusion_model_cfg) + + # ------------------------------------------------------------------ + # Action horizon (chunk length sent to the DiT) + # Single source of truth: `action_horizon` (e.g. 8). + # Legacy YAMLs that only provide `future_action_window_size` are + # normalised to `action_horizon` upstream by + # `share_tools.apply_config_compat`, so this code never touches + # the legacy alias. + # ------------------------------------------------------------------ + self.action_horizon = int(config.action_horizon) + + # ------------------------------------------------------------------ + # Action / state dimensions + # action_dim: DoF of the robot action (e.g. 7 for 6-DoF + gripper) + # state_dim: proprioception dimension; set to 0/None to disable + # the state_encoder branch entirely. + # ------------------------------------------------------------------ + self.action_dim = config.action_dim + + # ------------------------------------------------------------------ + # Inference denoising steps + # num_inference_timesteps: Euler steps during predict_action(). + # Typically 4–10; fewer = faster but less accurate. + # ------------------------------------------------------------------ + self.num_inference_timesteps = config.num_inference_timesteps + + # ------------------------------------------------------------------ + # hidden_size: intermediate MLP width for state_encoder / action_decoder. + # Decoupled from input_embedding_dim so you can use a smaller hidden + # for the MLP without changing the DiT latent size. + # ------------------------------------------------------------------ + self.hidden_size = config.hidden_size + + self.state_encoder = ( + MLP( + input_dim=config.state_dim, + hidden_dim=self.hidden_size, + output_dim=self.input_embedding_dim, + ) + if config.state_dim + else None + ) + + self.action_encoder = ActionEncoder( + action_dim=config.action_dim, + hidden_size=self.input_embedding_dim, + ) + self.action_decoder = MLP( + input_dim=self.model.config.output_dim, + hidden_dim=self.hidden_size, + output_dim=self.action_dim, + ) + + # ------------------------------------------------------------------ + # future_tokens: learnable query tokens prepended before the action + # sequence so the DiT has dedicated "planning" slots. + # num_target_vision_tokens controls how many such tokens are added. + # ------------------------------------------------------------------ + self.future_tokens = nn.Embedding(config.num_target_vision_tokens, self.input_embedding_dim) + nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02) + + # ------------------------------------------------------------------ + # Positional embedding over the action sequence + # add_pos_embed: whether to add sinusoidal-style learned PE + # max_seq_len: max supported action sequence length + # ------------------------------------------------------------------ + if config.add_pos_embed: + self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim) + nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02) + + # ------------------------------------------------------------------ + # Flow-matching noise schedule (Beta distribution) + # noise_beta_alpha / noise_beta_beta: Beta(α, β) shape params. + # noise_s: upper-clip of the sampled value so t ∈ [0, noise_s]. + # num_timestep_buckets: discretise continuous t into N buckets for + # the timestep encoder inside DiT. + # ------------------------------------------------------------------ + self.beta_dist = Beta(config.noise_beta_alpha, config.noise_beta_beta) + self.num_timestep_buckets = config.num_timestep_buckets + self.config = config + + def sample_time(self, batch_size, device, dtype): + sample = self.beta_dist.sample([batch_size]).to(device, dtype=dtype).clamp(max=self.config.noise_s) + return (self.config.noise_s - sample) / self.config.noise_s + + def _sample_training_noise_and_time( + self, + actions: torch.Tensor, + flow_noise_seed: int | None, + flow_noise: torch.Tensor | None = None, + flow_time: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if (flow_noise is None) != (flow_time is None): + raise ValueError("flow_noise and flow_time must be supplied together") + if flow_noise is not None: + if tuple(flow_noise.shape) != tuple(actions.shape): + raise ValueError("Explicit flow noise does not match the action shape") + if tuple(flow_time.shape) != (actions.shape[0],): + raise ValueError("Explicit flow time does not match the batch shape") + return ( + flow_noise.to(device=actions.device, dtype=actions.dtype), + flow_time.to(device=actions.device, dtype=actions.dtype), + ) + + def sample() -> tuple[torch.Tensor, torch.Tensor]: + noise = torch.randn( + actions.shape, + device=actions.device, + dtype=actions.dtype, + ) + time = self.sample_time( + actions.shape[0], + device=actions.device, + dtype=actions.dtype, + ) + return noise, time + + if flow_noise_seed is None: + return sample() + devices = ( + [torch.cuda.current_device()] + if actions.is_cuda + else [] + ) + # Flow draws must not advance the dropout stream. + with torch.random.fork_rng(devices=devices): + torch.manual_seed(int(flow_noise_seed)) + if actions.is_cuda: + torch.cuda.manual_seed(int(flow_noise_seed)) + return sample() + + def prepare_input(self, batch: dict) -> BatchFeature: + return BatchFeature(data=batch) + + def predict_velocity( + self, + vl_embs: torch.Tensor, + actions: torch.Tensor, + time: torch.Tensor, + state: torch.Tensor = None, + encoder_attention_mask=None, + ) -> torch.Tensor: + """Evaluate the conditional action vector field at continuous time.""" + if time.shape != (actions.shape[0],): + raise ValueError( + f"time must have shape {(actions.shape[0],)}, got {tuple(time.shape)}" + ) + time_float = time.to(device=actions.device, dtype=torch.float32).clamp( + min=0.0, max=1.0 + ) + timesteps = (time_float * self.num_timestep_buckets).floor().long() + timesteps = timesteps.clamp(max=self.num_timestep_buckets) + action_features = self.action_encoder(actions, timesteps) + + state_features = self.state_encoder(state) if state is not None else None + if self.config.add_pos_embed: + pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=actions.device) + action_features = action_features + self.position_embedding(pos_ids).unsqueeze(0) + + future_tokens = self.future_tokens.weight.unsqueeze(0).expand(vl_embs.shape[0], -1, -1) + sa_embs = ( + torch.cat((state_features, future_tokens, action_features), dim=1) + if state_features is not None + else torch.cat((future_tokens, action_features), dim=1) + ) + model_output = self.model( + hidden_states=sa_embs, + encoder_hidden_states=vl_embs, + encoder_attention_mask=encoder_attention_mask, + timestep=timesteps, + return_all_hidden_states=False, + ) + prediction = self.action_decoder(model_output) + return prediction[:, -actions.shape[1] :] + + def forward( + self, + vl_embs: torch.Tensor, + actions: torch.Tensor, + state: torch.Tensor = None, + encoder_attention_mask=None, + reduction: str = "mean", + flow_noise_seed: int | None = None, + flow_noise: torch.Tensor | None = None, + flow_time: torch.Tensor | None = None, + return_flow_noise_time: bool = False, + ): + """ + vl_embs: shape (B, seq_length, feature_dim) + actions: shape (B, action_horizon, action_dim) + """ + # Embed noised action trajectory. + noise, t = self._sample_training_noise_and_time( + actions, + flow_noise_seed, + flow_noise=flow_noise, + flow_time=flow_time, + ) + t = t[:, None, None] # shape (B,1,1) for broadcast + + noisy_trajectory = (1 - t) * noise + t * actions + velocity = actions - noise + + pred_actions = self.predict_velocity( + vl_embs, + noisy_trajectory, + t[:, 0, 0], + state=state, + encoder_attention_mask=encoder_attention_mask, + ) + + # Slice out only the action portion of pred and target. + elementwise_loss = (pred_actions - velocity) ** 2 + loss = reduce_flow_matching_loss(elementwise_loss, reduction=reduction) + if return_flow_noise_time: + return loss, noise, t[:, 0, 0] + return loss + + @torch.no_grad() + def predict_action( + self, + vl_embs: torch.Tensor, + state: torch.Tensor = None, + encoder_attention_mask=None, + ) -> torch.Tensor: + # Set initial actions as the sampled noise. + batch_size = vl_embs.shape[0] + device = vl_embs.device + actions = torch.randn( + size=(batch_size, self.action_horizon, self.action_dim), + dtype=vl_embs.dtype, + device=device, + ) + + num_steps = self.num_inference_timesteps + dt = 1.0 / num_steps + + # Run denoising steps. + for t in range(num_steps): + t_cont = t / float(num_steps) # e.g. goes 0, 1/N, 2/N, ... + times = torch.full( + (batch_size,), t_cont, device=device, dtype=torch.float32 + ) + pred_velocity = self.predict_velocity( + vl_embs, + actions, + times, + state=state, + encoder_attention_mask=encoder_attention_mask, + ) + + # Update actions using euler integration. + actions = actions + dt * pred_velocity + return actions + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + +def get_action_model(config=None): + """ + Factory: build FlowmatchingActionHead from global framework config. + + Args: + config: Global config (expects config.framework.action_model namespace). + + Returns: + FlowmatchingActionHead: Initialized FlowMatchingActionHead. + """ + return FlowmatchingActionHead(full_config=config) + + +if __name__ == "__main__": + # TODO make each backbone.py can be debug independently + + pass diff --git a/real_robot/source/starVLA/model/modules/action_model/__init__.py b/real_robot/source/starVLA/model/modules/action_model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b09bb11c9b2bf0ed931612e81f6f729cddd0c162 --- /dev/null +++ b/real_robot/source/starVLA/model/modules/action_model/__init__.py @@ -0,0 +1,2 @@ +# Modified in the ReTime research artifact; see NOTICE. +"""Action models included in the Florence-Bridge OPTR release.""" diff --git a/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/__init__.py b/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3159bfe65645499015bd92609b99d476d69544e9 --- /dev/null +++ b/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py b/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..e767827402098ca666115195f66cccf6890b4945 --- /dev/null +++ b/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import torch +import torch.nn as nn + + +def swish(x): + return x * torch.sigmoid(x) + + +class SinusoidalPositionalEncoding(nn.Module): + """ + Produces a sinusoidal encoding of shape (B, T, w) + given timesteps of shape (B, T). + """ + + def __init__(self, embedding_dim): + super().__init__() + self.embedding_dim = embedding_dim + + def forward(self, timesteps): + # timesteps: shape (B, T) + # We'll compute sin/cos frequencies across dim T + timesteps = timesteps.float() # ensure float + + B, T = timesteps.shape + device = timesteps.device + + half_dim = self.embedding_dim // 2 + # typical log space frequencies for sinusoidal encoding + exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( + torch.log(torch.tensor(10000.0)) / half_dim + ) + # Expand timesteps to (B, T, 1) then multiply + freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) + + sin = torch.sin(freqs) + cos = torch.cos(freqs) + enc = torch.cat([sin, cos], dim=-1) # (B, T, w) + + return enc + + +class ActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size): + super().__init__() + self.hidden_size = hidden_size + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = nn.Linear(action_dim, hidden_size) # (d -> w) + self.W2 = nn.Linear(2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = nn.Linear(hidden_size, hidden_size) # (w -> w) + + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x) + + return x diff --git a/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py b/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..412dd8a7f8a6fe08d799c4b1f8e38176a8eb2ed9 --- /dev/null +++ b/real_robot/source/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py @@ -0,0 +1,380 @@ +# Modified in the ReTime research artifact; see NOTICE. +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 Optional + +import torch +import torch.nn.functional as F +from diffusers import ConfigMixin, ModelMixin +from diffusers.configuration_utils import register_to_config +from diffusers.models.attention import Attention, FeedForward +from diffusers.models.embeddings import ( + SinusoidalPositionalEmbedding, + TimestepEmbedding, + Timesteps, +) +from torch import nn + + +class TimestepEncoder(nn.Module): + def __init__(self, embedding_dim, compute_dtype=torch.float32): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timesteps): + dtype = next(self.parameters()).dtype + timesteps_proj = self.time_proj(timesteps).to(dtype) + timesteps_emb = self.timestep_embedder(timesteps_proj) # (N, D) + return timesteps_emb + + +class AdaLayerNorm(nn.Module): + def __init__( + self, + embedding_dim: int, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + chunk_dim: int = 0, + ): + super().__init__() + self.chunk_dim = chunk_dim + output_dim = embedding_dim * 2 + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, output_dim) + self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine) + + def forward( + self, + x: torch.Tensor, + temb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + temb = self.linear(self.silu(temb)) + scale, shift = temb.chunk(2, dim=1) + x = self.norm(x) * (1 + scale[:, None]) + shift[:, None] + return x + + +class BasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + attention_bias: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = True, + attention_out_bias: bool = True, + ): + super().__init__() + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + self.dropout = dropout + self.cross_attention_dim = cross_attention_dim + self.activation_fn = activation_fn + self.attention_bias = attention_bias + self.norm_elementwise_affine = norm_elementwise_affine + self.positional_embeddings = positional_embeddings + self.num_positional_embeddings = num_positional_embeddings + self.norm_type = norm_type + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if norm_type == "ada_norm": + self.norm1 = AdaLayerNorm(dim) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim, + upcast_attention=upcast_attention, + out_bias=attention_out_bias, + ) + + # 3. Feed-forward + self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + inner_dim=ff_inner_dim, + bias=ff_bias, + ) + if final_dropout: + self.final_dropout = nn.Dropout(dropout) + else: + self.final_dropout = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + temb: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + + # 0. Self-Attention + if self.norm_type == "ada_norm": + norm_hidden_states = self.norm1(hidden_states, temb) + else: + norm_hidden_states = self.norm1(hidden_states) + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + ) + if self.final_dropout: + attn_output = self.final_dropout(attn_output) + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 4. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + ff_output = self.ff(norm_hidden_states) + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + return hidden_states + + +class DiT(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + # register_to_config auto-registers constructor params into config, enabling access via self.config.xxx instead of self.xxx + @register_to_config # Registers passed params to config. TODO: replace with our singleton pattern, implement a mergeable @merge_param_config + def __init__( + self, + num_attention_heads: int = 8, + attention_head_dim: int = 64, + output_dim: int = 26, + num_layers: int = 12, + dropout: float = 0.1, + attention_bias: bool = True, + activation_fn: str = "gelu-approximate", + num_embeds_ada_norm: Optional[int] = 1000, + upcast_attention: bool = False, + norm_type: str = "ada_norm", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + max_num_positional_embeddings: int = 512, + compute_dtype=torch.float32, + final_dropout: bool = True, + positional_embeddings: Optional[str] = "sinusoidal", + interleave_self_attention=False, + cross_attention_dim: Optional[int] = None, + **kwargs, + ): + super().__init__() + self.attention_head_dim = attention_head_dim + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.gradient_checkpointing = False + + # Timestep encoder + # self.config.compute_dtype may not exist, handle it in advance + compute_dtype = getattr(self.config, "compute_dtype", torch.float32) + self.timestep_encoder = ( + TimestepEncoder( # TODO BUG: self.config.compute_dtype doesn't error during training but fails at eval + embedding_dim=self.inner_dim, compute_dtype=compute_dtype + ) + ) + + all_blocks = [] + for idx in range(self.config.num_layers): + + use_self_attn = idx % 2 == 1 and interleave_self_attention + curr_cross_attention_dim = cross_attention_dim if not use_self_attn else None + + all_blocks += [ + BasicTransformerBlock( + self.inner_dim, + self.config.num_attention_heads, + self.config.attention_head_dim, + dropout=self.config.dropout, + activation_fn=self.config.activation_fn, + attention_bias=self.config.attention_bias, + upcast_attention=self.config.upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=self.config.norm_elementwise_affine, + norm_eps=self.config.norm_eps, + positional_embeddings=positional_embeddings, + num_positional_embeddings=self.config.max_num_positional_embeddings, + final_dropout=final_dropout, + cross_attention_dim=curr_cross_attention_dim, + ) + ] + self.transformer_blocks = nn.ModuleList(all_blocks) + + # Output blocks + self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) + self.proj_out_2 = nn.Linear(self.inner_dim, self.config.output_dim) + print( + "Total number of DiT parameters: ", + sum(p.numel() for p in self.parameters() if p.requires_grad), + ) + + def forward( + self, + hidden_states: torch.Tensor, # Shape: (B, T, D) + encoder_hidden_states: torch.Tensor, # Shape: (B, S, D) + timestep: Optional[torch.LongTensor] = None, + return_all_hidden_states: bool = False, + encoder_attention_mask=None, + ): + # Encode timesteps + temb = self.timestep_encoder(timestep) + + # Process through transformer blocks - single pass through the blocks + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + + all_hidden_states = [hidden_states] if return_all_hidden_states else None + + # Process through transformer blocks + for idx, block in enumerate(self.transformer_blocks): + if idx % 2 == 1 and self.config.interleave_self_attention: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + temb=temb, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=None, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + temb=temb, + ) + if all_hidden_states is not None: + all_hidden_states.append(hidden_states) + + # Output processing + conditioning = temb + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + if return_all_hidden_states: + return self.proj_out_2(hidden_states), all_hidden_states + else: + return self.proj_out_2(hidden_states) + + +class SelfAttentionTransformer(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 8, + attention_head_dim: int = 64, + output_dim: int = 26, + num_layers: int = 12, + dropout: float = 0.1, + attention_bias: bool = True, + activation_fn: str = "gelu-approximate", + num_embeds_ada_norm: Optional[int] = 1000, + upcast_attention: bool = False, + max_num_positional_embeddings: int = 512, + compute_dtype=torch.float32, + final_dropout: bool = True, + positional_embeddings: Optional[str] = "sinusoidal", + interleave_self_attention=False, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.gradient_checkpointing = False + + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + self.inner_dim, + self.config.num_attention_heads, + self.config.attention_head_dim, + dropout=self.config.dropout, + activation_fn=self.config.activation_fn, + attention_bias=self.config.attention_bias, + upcast_attention=self.config.upcast_attention, + positional_embeddings=positional_embeddings, + num_positional_embeddings=self.config.max_num_positional_embeddings, + final_dropout=final_dropout, + ) + for _ in range(self.config.num_layers) + ] + ) + print( + "Total number of SelfAttentionTransformer parameters: ", + sum(p.numel() for p in self.parameters() if p.requires_grad), + ) + + def forward( + self, + hidden_states: torch.Tensor, # Shape: (B, T, D) + return_all_hidden_states: bool = False, + ): + # Process through transformer blocks - single pass through the blocks + hidden_states = hidden_states.contiguous() + all_hidden_states = [hidden_states] if return_all_hidden_states else None + + # Process through transformer blocks + for idx, block in enumerate(self.transformer_blocks): + hidden_states = block(hidden_states) + if all_hidden_states is not None: + all_hidden_states.append(hidden_states) + + if return_all_hidden_states: + return hidden_states, all_hidden_states + else: + return hidden_states diff --git a/real_robot/source/starVLA/model/modules/vlm/Florence2.py b/real_robot/source/starVLA/model/modules/vlm/Florence2.py new file mode 100644 index 0000000000000000000000000000000000000000..a937144ab85a43ecd300baeac3ca37dd5dac6447 --- /dev/null +++ b/real_robot/source/starVLA/model/modules/vlm/Florence2.py @@ -0,0 +1,340 @@ +# Modified in the ReTime research artifact; see NOTICE. +# Copyright 2025 starVLA community. All rights reserved. +# Licensed under the MIT License, Version 1.0 (the "License"); + +from typing import Optional + +import torch +from starVLA.training.trainer_utils import initialize_overwatch +from PIL import Image +from transformers import AutoModelForCausalLM, AutoProcessor +from transformers.modeling_outputs import CausalLMOutputWithPast + +logger = initialize_overwatch(__name__) + +# IGNORE_INDEX = -100 +# IMAGE_TOKEN_INDEX = 151655 +# VIDEO_TOKEN_INDEX = 151656 +# DEFAULT_IMAGE_TOKEN = "" +# DEFAULT_VIDEO_TOKEN = "